diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py
index 9dd4aff3dc..8ed18bfd12 100644
--- a/src/maxtext/common/common_types.py
+++ b/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):
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 55925ffd58..1091ff7d8d 100644
--- a/src/maxtext/configs/base.yml
+++ b/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
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 03640ec758..ccc817b090 100644
--- a/src/maxtext/configs/types.py
+++ b/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 (
diff --git a/src/maxtext/diffusion/__init__.py b/src/maxtext/diffusion/__init__.py
new file mode 100644
index 0000000000..2c9fc63d21
--- /dev/null
+++ b/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."""
diff --git a/src/maxtext/diffusion/denoise.py b/src/maxtext/diffusion/denoise.py
new file mode 100644
index 0000000000..2f99b91069
--- /dev/null
+++ b/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)
diff --git a/src/maxtext/diffusion/scoring.py b/src/maxtext/diffusion/scoring.py
new file mode 100644
index 0000000000..bf220219dc
--- /dev/null
+++ b/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}")
diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py
index 370f1895bd..a047baadbb 100644
--- a/src/maxtext/input_pipeline/hf_data_processing.py
+++ b/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.
diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py
index d996a88908..6b294a6dcb 100644
--- a/src/maxtext/input_pipeline/input_pipeline_utils.py
+++ b/src/maxtext/input_pipeline/input_pipeline_utils.py
@@ -325,14 +325,25 @@ def tokenization(example, hf_tokenizer, truncation, max_length, column_names):
 @dataclasses.dataclass
 class SFTPromptMasking(grain.MapTransform):
   """Construct inputs and targets for SFT training. Concat prompt and completion to generate inputs.
-  For targets, if train on completion only, the prompt will be masked by unk_id. Otherwise the same as inputs.
+  Causal SFT can mask prompt targets with ``unk_id``. Target-aligned objectives
+  keep clean targets and may emit role-derived completion metadata instead.
   """
 
-  def __init__(self, text_column_name, completion_only, max_target_length, unk_id=0):
+  def __init__(
+      self,
+      text_column_name,
+      completion_only,
+      max_target_length,
+      unk_id=0,
+      target_aligned=False,
+      emit_completion_mask=False,
+  ):
     self.text_column_name = text_column_name
     self.completion_only = completion_only
     self.max_target_length = max_target_length
     self.unk_id = unk_id
+    self.target_aligned = target_aligned
+    self.emit_completion_mask = emit_completion_mask
 
   def map(self, element):
     """
@@ -340,18 +351,24 @@ def map(self, element):
     It concatenates the prompt and completion to form the `inputs` sequence.
     For the `targets` sequence:
     - If `self.completion_only` is `True`, the prompt portion of the
-      concatenated sequence is masked using `self.unk_id`.
+      concatenated sequence is masked using `self.unk_id`, unless
+      `self.target_aligned` requests clean same-position targets.
     - If `self.completion_only` is `False`, the target sequence is
       identical to the input sequence.
     """
-    inputs, targets = [], []
+    inputs, targets, completion_mask = [], [], []
     for i, text in enumerate(element[self.text_column_name]):
       inputs += text
-      targets += [self.unk_id] * len(text) if self.completion_only and element["is_prompt"][i] else text
-    return {
+      is_prompt = element["is_prompt"][i]
+      targets += [self.unk_id] * len(text) if self.completion_only and is_prompt and not self.target_aligned else text
+      completion_mask += [not is_prompt] * len(text)
+    output = {
         "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32),
         "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32),
     }
+    if self.emit_completion_mask:
+      output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32)
+    return output
 
 
 @dataclasses.dataclass
@@ -838,7 +855,7 @@ def map(
     """map to each element"""
     data_columns = list(element.keys())
     for data_column in data_columns:
-      if data_column != "images":
+      if data_column not in ("images", "completion_mask"):
         if isinstance(element[data_column], mm_utils.PreprocessorOutput):
           raise TypeError("Only 'images' column can be of type PreprocessorOutput.")
 
@@ -863,6 +880,8 @@ def map(
 
         element["images"] = self._pad_image_and_mask(element["images"])  # pyrefly: ignore[bad-argument-type]
 
+      elif key == "completion_mask":
+        element[key] = self._pad_text(element[key], self.max_length, 0)  # pyrefly: ignore[bad-argument-type]
       elif "true_length" not in key:
         element[key] = self._pad_text(element[key], self.max_length, self.pad_id)  # pyrefly: ignore[bad-argument-type]
     return element
@@ -986,6 +1005,90 @@ def map(self, element):
     return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis)
 
 
+@dataclasses.dataclass
+class BlockDiffusionCorruption(grain.RandomMapTransform):
+  """Builds explicit completion, corruption, and loss masks for block diffusion."""
+
+  def __init__(
+      self,
+      block_size: int,
+      mask_id: int,
+      min_noise: float = 1.0e-3,
+      axis: int = 1,
+      completion_only: bool = False,
+      seed_first_token: bool = False,
+      include_seed_in_loss: bool = False,
+  ):
+    if block_size <= 0:
+      raise ValueError(f"block_size must be positive, got {block_size}")
+    if not 0.0 < min_noise <= 1.0:
+      raise ValueError(f"min_noise must satisfy 0 < min_noise <= 1, got {min_noise}")
+    if include_seed_in_loss and not seed_first_token:
+      raise ValueError("include_seed_in_loss requires seed_first_token=True")
+    self.block_size = block_size
+    self.mask_id = mask_id
+    self.min_noise = min_noise
+    self.axis = axis
+    self.completion_only = completion_only
+    self.seed_first_token = seed_first_token
+    self.include_seed_in_loss = include_seed_in_loss
+
+  def random_map(self, element, rng: np.random.Generator):
+    inputs = np.asarray(element["inputs"])
+    targets = np.asarray(element["targets"])
+    targets_segmentation = np.asarray(element["targets_segmentation"])
+    if inputs.shape != targets.shape or inputs.shape != targets_segmentation.shape:
+      raise ValueError(
+          "inputs, targets, and targets_segmentation must have identical shapes, got "
+          f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}"
+      )
+    if inputs.ndim == 0 or not -inputs.ndim <= self.axis < inputs.ndim:
+      raise ValueError(f"axis {self.axis} is invalid for inputs with {inputs.ndim} dimensions")
+
+    axis = self.axis % inputs.ndim
+    validity_mask = targets_segmentation != 0
+    if self.completion_only and "completion_mask" not in element:
+      raise ValueError("completion_only block diffusion requires an explicit completion_mask")
+    completion_mask = np.asarray(element.get("completion_mask", validity_mask)) != 0
+    if completion_mask.shape != inputs.shape:
+      raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}")
+    completion_mask &= validity_mask
+    eligible_mask = completion_mask if self.completion_only else validity_mask
+
+    moved_eligible = np.moveaxis(eligible_mask, axis, -1)
+    rows = moved_eligible.reshape(-1, inputs.shape[axis])
+    block_count = (rows.shape[-1] + self.block_size - 1) // self.block_size
+    padded_length = block_count * self.block_size
+    padded_rows = np.pad(rows, ((0, 0), (0, padded_length - rows.shape[-1])))
+    eligible_blocks = padded_rows.reshape(rows.shape[0], block_count, self.block_size)
+    seed_loss_blocks = np.zeros_like(eligible_blocks)
+    if self.seed_first_token:
+      if self.include_seed_in_loss:
+        seed_loss_blocks = eligible_blocks & (np.arange(self.block_size) == 0)
+        seed_loss_blocks &= np.arange(block_count)[None, :, None] != 0
+      eligible_blocks = eligible_blocks & (np.arange(self.block_size) != 0)
+
+    noise = rng.uniform(self.min_noise, 1.0, size=eligible_blocks.shape[:-1] + (1,))
+    selected_blocks = eligible_blocks & (rng.random(size=eligible_blocks.shape) < noise)
+    needs_fallback = eligible_blocks.any(axis=-1) & ~selected_blocks.any(axis=-1)
+    fallback_scores = np.where(eligible_blocks, rng.random(size=eligible_blocks.shape), -1.0)
+    fallback_offsets = np.argmax(fallback_scores, axis=-1)
+    fallback = np.arange(self.block_size) == fallback_offsets[..., None]
+    selected_blocks |= needs_fallback[..., None] & fallback
+
+    corruption_rows = selected_blocks.reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]]
+    loss_rows = (selected_blocks | seed_loss_blocks).reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]]
+    corruption_mask = np.moveaxis(corruption_rows.reshape(moved_eligible.shape), -1, axis)
+    targets_loss_mask = np.moveaxis(loss_rows.reshape(moved_eligible.shape), -1, axis)
+    output = dict(element)
+    output["inputs"] = np.where(corruption_mask, inputs.dtype.type(self.mask_id), inputs)
+    output["targets"] = targets
+    output["completion_mask"] = completion_mask.astype(targets_segmentation.dtype)
+    output["corruption_mask"] = corruption_mask.astype(targets_segmentation.dtype)
+    output["targets_loss_mask"] = targets_loss_mask.astype(targets_segmentation.dtype)
+    return output
+
+
 @dataclasses.dataclass
 class ComputeQwen3OmniPositions(grain.MapTransform):
   """Computes 3D position IDs for Qwen3-Omni multimodal sequences.
diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py
index d10816435c..a184e522cd 100644
--- a/src/maxtext/input_pipeline/synthetic_data_processing.py
+++ b/src/maxtext/input_pipeline/synthetic_data_processing.py
@@ -111,6 +111,11 @@ def get_place_holder_synthetic_data(config: pyconfig.HyperParameters):
         "targets_position": neg_ones,
         "targets_segmentation": neg_ones,
     }
+    if getattr(config, "training_objective", "causal_lm") == "block_diffusion":
+      zeros = np.zeros_like(neg_ones)
+      batch["completion_mask"] = zeros
+      batch["corruption_mask"] = zeros
+      batch["targets_loss_mask"] = zeros
 
     def infinite_iterator():
       while True:
diff --git a/src/maxtext/integration/tunix/diffusion_sft.py b/src/maxtext/integration/tunix/diffusion_sft.py
new file mode 100644
index 0000000000..b85132db38
--- /dev/null
+++ b/src/maxtext/integration/tunix/diffusion_sft.py
@@ -0,0 +1,180 @@
+# 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.
+
+"""MaxText adapters for target-aligned Tunix diffusion SFT."""
+
+from collections.abc import Mapping
+from typing import Any
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.diffusion import scoring
+from tunix.diffusion import types as diffusion_types
+
+
+_REQUIRED_FIELDS = (
+    "inputs",
+    "inputs_position",
+    "inputs_segmentation",
+    "targets",
+    "targets_position",
+    "targets_segmentation",
+    "completion_mask",
+    "corruption_mask",
+    "targets_loss_mask",
+)
+
+
+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_batch_masks(
+    positions,
+    validity_mask,
+    completion_mask,
+    corruption_mask,
+    loss_weights,
+    *,
+    alignment,
+    block_size,
+):
+    shapes = {
+        "positions": tuple(positions.shape),
+        "validity_mask": tuple(validity_mask.shape),
+        "completion_mask": tuple(completion_mask.shape),
+        "corruption_mask": tuple(corruption_mask.shape),
+        "targets_loss_mask": tuple(loss_weights.shape),
+    }
+    if len(set(shapes.values())) != 1:
+        raise ValueError(
+            f"diffusion SFT masks must have identical shapes; received {shapes}"
+        )
+    concrete = [
+        _concrete_numpy(value)
+        for value in (
+            positions,
+            validity_mask,
+            completion_mask,
+            corruption_mask,
+            loss_weights,
+        )
+    ]
+    if any(value is None for value in concrete):
+        return
+    concrete_positions = np.asarray(concrete[0])
+    validity, completion, corruption = (
+        np.asarray(value, dtype=bool) for value in concrete[1:4]
+    )
+    weights = np.asarray(concrete[4])
+    weighted = weights != 0
+    if not np.all(np.isfinite(weights)) or np.any(weights < 0):
+        raise ValueError("targets_loss_mask must contain finite nonnegative weights")
+    if np.any(completion & ~validity):
+        raise ValueError("completion_mask must be a subset of valid target positions")
+    if np.any(corruption & ~completion):
+        raise ValueError("corruption_mask must be a subset of completion_mask")
+    if np.any(weighted & ~completion):
+        raise ValueError("targets_loss_mask must be a subset of completion_mask")
+    allowed = corruption.copy()
+    if alignment == "shifted":
+        allowed |= (concrete_positions > 0) & (concrete_positions % block_size == 0)
+    if np.any(weighted & ~allowed):
+        raise ValueError(
+            "diffusion SFT loss weights must own corrupted targets or shifted block anchors"
+        )
+
+
+def create_batch_adapter(config):
+    """Builds a raw MaxText batch adapter for Tunix diffusion SFT."""
+    alignment = config.block_diffusion_logit_alignment
+    block_size = int(config.block_diffusion_block_size)
+
+    def adapt(raw_batch: Mapping[str, Any]) -> diffusion_types.DiffusionTokenBatch:
+        missing = sorted(name for name in _REQUIRED_FIELDS if name not in raw_batch)
+        if missing:
+            raise ValueError(
+                f"block-diffusion SFT requires explicit batch fields; missing {missing}"
+            )
+        targets = jnp.asarray(raw_batch["targets"])
+        positions = jnp.asarray(raw_batch["targets_position"], dtype=jnp.int32)
+        validity_mask = jnp.asarray(raw_batch["targets_segmentation"]) != 0
+        completion_mask = jnp.asarray(raw_batch["completion_mask"], dtype=jnp.bool_)
+        corruption_mask = jnp.asarray(raw_batch["corruption_mask"], dtype=jnp.bool_)
+        raw_loss_weights = jnp.asarray(
+            raw_batch["targets_loss_mask"], dtype=jnp.float32
+        )
+        _validate_batch_masks(
+            positions,
+            validity_mask,
+            completion_mask,
+            corruption_mask,
+            raw_loss_weights,
+            alignment=alignment,
+            block_size=block_size,
+        )
+        allowed = corruption_mask
+        if alignment == "shifted":
+            allowed |= (positions > 0) & (positions % block_size == 0)
+        loss_weights = jnp.where(
+            validity_mask & completion_mask & allowed, raw_loss_weights, 0.0
+        )
+        return diffusion_types.DiffusionTokenBatch.create(
+            model_inputs={
+                "input_tokens": jnp.asarray(raw_batch["inputs"]),
+                "input_positions": jnp.asarray(
+                    raw_batch["inputs_position"], dtype=jnp.int32
+                ),
+                "input_segmentation": jnp.asarray(raw_batch["inputs_segmentation"]),
+                "targets": targets,
+                "target_positions": positions,
+                "target_segmentation": jnp.asarray(raw_batch["targets_segmentation"]),
+            },
+            target_ids=targets,
+            loss_weights=loss_weights,
+        )
+
+    return adapt
+
+
+def create_target_aligned_logits_fn(config):
+    """Builds a MaxText scorer satisfying Tunix's diffusion logits contract."""
+    alignment = config.block_diffusion_logit_alignment
+    enable_dropout = bool(config.enable_dropout)
+
+    def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs):
+        base_model = getattr(model, "base", model)
+        logits = base_model(
+            decoder_input_tokens=model_inputs["input_tokens"],
+            decoder_positions=model_inputs["input_positions"],
+            decoder_segment_ids=model_inputs["input_segmentation"],
+            enable_dropout=enable_dropout,
+            decoder_target_tokens=model_inputs["targets"],
+            decoder_target_mask=model_inputs["target_segmentation"],
+        )
+        return scoring.align_logits_to_targets(
+            logits,
+            alignment,
+            model_inputs["target_positions"],
+            model_inputs["target_segmentation"] != 0,
+        )
+
+    return logits_fn
diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py
index 8e6c96a72e..af48802cc9 100644
--- a/src/maxtext/layers/attention_op.py
+++ b/src/maxtext/layers/attention_op.py
@@ -84,6 +84,16 @@
 dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None))
 
 
+def _resolve_attention_type(config: Config, attention_type: AttentionType | str | None) -> AttentionType:
+  configured_attention_type = AttentionType(getattr(config, "attention_type", AttentionType.GLOBAL.value))
+  if attention_type is None:
+    return configured_attention_type
+  resolved_attention_type = AttentionType(attention_type)
+  if configured_attention_type == AttentionType.BLOCK_DIFFUSION and resolved_attention_type == AttentionType.GLOBAL:
+    return configured_attention_type
+  return resolved_attention_type
+
+
 def validate_compute_axis_order(s: AxisIdxes) -> None:
   valid_compute_axis_order = ((0, 1, 2, 3), (0, 2, 1, 3))
   if s not in valid_compute_axis_order:  # currently supported compute_axis_order
@@ -204,6 +214,50 @@ def __hash__(self):
     )
 
 
+class BlockCausalMask(splash_attention_mask._ComputableMask):  # pylint: disable=protected-access,abstract-method
+  """Lazy mask with bidirectional attention within causal blocks."""
+
+  block_size: int
+
+  def __init__(
+      self,
+      shape: tuple[int, int],
+      block_size: int,
+      shard_count: int = 1,
+  ):
+    if block_size <= 0:
+      raise ValueError("block_size must be positive")
+    self.block_size = block_size
+
+    def block_causal_mask_function(q_ids, kv_ids):
+      return (q_ids // self.block_size) >= (kv_ids // self.block_size)
+
+    super().__init__(
+        shape=shape,
+        mask_function=block_causal_mask_function,
+        shard_count=shard_count,
+    )
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+    return (
+        self.shape == other.shape
+        and self.block_size == other.block_size
+        and np.array_equal(self.q_sequence, other.q_sequence)
+    )
+
+  def __hash__(self):
+    return hash(
+        (
+            type(self),
+            self.shape,
+            self.block_size,
+            self.q_sequence.tobytes() if self.q_sequence is not None else None,
+        )
+    )
+
+
 def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, q_offset: int = 0) -> jax.Array:
   """Generates an explicit boolean mask for chunked causal attention.
 
@@ -234,6 +288,16 @@ def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int,
   return chunk_mask
 
 
+def _generate_block_causal_attention_mask(mask_shape: tuple[int, int], block_size: int, q_offset: int = 0) -> jax.Array:
+  """Generates a block-causal mask, bidirectional within each block."""
+  if block_size <= 0:
+    raise ValueError("block_size must be positive")
+
+  row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + q_offset
+  col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+  return (row_ids // block_size) >= (col_ids // block_size)
+
+
 def _make_block_mask_indices(bidirectional_mask):
   """Creates block mask identifying segments based on a bidirectional mask.
 
@@ -479,7 +543,13 @@ def __init__(
     self.dtype = dtype
     self.quant = quant
     self.kv_quant = kv_quant
-    self.attention_type = attention_type
+    self.attention_type = _resolve_attention_type(self.config, attention_type)
+    self.block_diffusion_block_size = getattr(self.config, "block_diffusion_block_size", None)
+    if self.attention_type == AttentionType.BLOCK_DIFFUSION:
+      if self.block_diffusion_block_size is None or self.block_diffusion_block_size <= 0:
+        raise ValueError("block_diffusion_block_size must be positive for block-diffusion attention")
+      if self.attention_kernel not in ("autoselected", "dot_product", "flash"):
+        raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.")
     # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels
     if self.attention_kernel not in (
         "dot_product",
@@ -782,6 +852,7 @@ def generate_attention_mask(
     if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in (
         AttentionType.FULL,
         AttentionType.COMPRESSED,
+        AttentionType.BLOCK_DIFFUSION,
     ):
       if use_segment_positions:
         causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :]
@@ -861,6 +932,19 @@ def generate_attention_mask(
         )
       output_mask = chunk_mask * output_mask
 
+    elif self.attention_type == AttentionType.BLOCK_DIFFUSION and model_mode != MODEL_MODE_AUTOREGRESSIVE:
+      if use_segment_positions:
+        block_mask = (
+            (position_row_ids // self.block_diffusion_block_size) >= (position_col_ids // self.block_diffusion_block_size)
+        )[:, None, None, :, :]
+      else:
+        block_mask = _generate_block_causal_attention_mask(
+            mask_shape=(q_seq_len, kv_seq_len),
+            block_size=self.block_diffusion_block_size,
+            q_offset=next_pos,
+        )[None, None, None, :, :]
+      output_mask = block_mask if output_mask is None else jnp.logical_and(output_mask, block_mask)
+
     if bidirectional_mask is not None:
       image_mask = _make_bidirectional_block_mask(bidirectional_mask)
       output_mask = output_mask | image_mask[:, None, None, ...]
@@ -1126,6 +1210,11 @@ def apply_attention(
         return out, None, None
 
       else:
+        if self.attention_type == AttentionType.BLOCK_DIFFUSION:
+          raise ValueError(
+              "Block-diffusion flash attention is supported only by TPU Splash; "
+              "use attention='dot_product' on other hardware."
+          )
         if model_mode == MODEL_MODE_AUTOREGRESSIVE:
           # fallback to dot_product as pallas gpu flash attention doesn't support decode stage
           return self.apply_attention_dot(
@@ -1430,13 +1519,24 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
       sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap)
       mask_shape = (query.shape[2], key.shape[2])  # (q_seq_len, kv_seq_len)
       mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask
+      use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel
       if self.attention_type == AttentionType.FULL:
         mask = mask_module.FullMask(mask_shape)
+      elif self.attention_type == AttentionType.BLOCK_DIFFUSION:
+        mask_type = LoadBalancedBlockCausalMask if use_load_balanced_cp else BlockCausalMask
+        mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {}
+        mask = mask_type(
+            shape=mask_shape,
+            block_size=self.block_diffusion_block_size,
+            **mask_kwargs,
+        )
       else:
         mask = mask_module.CausalMask(shape=mask_shape)
 
-      use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel
-      if use_load_balanced_cp and self.attention_type != AttentionType.FULL:
+      if use_load_balanced_cp and self.attention_type not in (
+          AttentionType.FULL,
+          AttentionType.BLOCK_DIFFUSION,
+      ):
         mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size)
 
       # Apply local masking if local sliding attention is enabled.
@@ -2447,3 +2547,21 @@ def __init__(
         shard_count=shard_count,
     )
     self.q_sequence = _load_balanced_q_sequence(shape, cp_size)
+
+
+class LoadBalancedBlockCausalMask(BlockCausalMask):  # pylint: disable=abstract-method
+  """Lazy block-causal mask with load-balanced query positions."""
+
+  def __init__(
+      self,
+      shape: tuple[int, int],
+      block_size: int,
+      cp_size: int,
+      shard_count: int = 1,
+  ):
+    super().__init__(
+        shape=shape,
+        block_size=block_size,
+        shard_count=shard_count,
+    )
+    self.q_sequence = _load_balanced_q_sequence(shape, cp_size)
diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py
index 5a65b35ba4..32f233cbeb 100644
--- a/src/maxtext/layers/attentions.py
+++ b/src/maxtext/layers/attentions.py
@@ -51,7 +51,7 @@
     AttentionType,
 )
 from maxtext.layers import nnx_wrappers
-from maxtext.layers.attention_op import AttentionOp
+from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type
 from maxtext.layers.embeddings import (
     LLaMARotaryEmbedding,
     LlamaVisionRotaryEmbedding,
@@ -120,7 +120,7 @@ def attention_as_linen(
     float32_logits: bool = False,  # cast logits in float32 for stability.
     quant: Optional[Quant] = None,
     kv_quant: Optional[KVQuant] = None,
-    attention_type: AttentionType = AttentionType.GLOBAL,  # Default to global attention
+    attention_type: AttentionType | None = None,
     attn_logits_soft_cap: float | None = None,
     sliding_window_size: int | None = None,
     use_ragged_attention: bool = False,
@@ -277,7 +277,7 @@ def __init__(
       float32_logits: bool = False,  # cast logits in float32 for stability.
       quant: Optional[Quant] = None,
       kv_quant: Optional[KVQuant] = None,
-      attention_type: AttentionType = AttentionType.GLOBAL,  # Default to global attention
+      attention_type: AttentionType | None = None,
       attn_logits_soft_cap: float | None = None,
       sliding_window_size: int | None = None,
       use_ragged_attention: bool = False,
@@ -388,7 +388,7 @@ def __init__(
     self.float32_logits = float32_logits
     self.quant = quant
     self.kv_quant = kv_quant
-    self.attention_type = attention_type
+    self.attention_type = _resolve_attention_type(self.config, attention_type)
     self.attn_logits_soft_cap = attn_logits_soft_cap
     self.sliding_window_size = sliding_window_size
     self.use_ragged_attention = use_ragged_attention
diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py
index 9f30a3d6d8..03edfbcdb9 100644
--- a/src/maxtext/trainers/post_train/sft/hooks.py
+++ b/src/maxtext/trainers/post_train/sft/hooks.py
@@ -29,7 +29,7 @@ class SFTTrainingHooks(BaseTrainingHooks):
   @override
   def get_total_weights(self, batch) -> jax.Array:
     """Calculate the number of non-padded tokens in the batch."""
-    return jnp.sum(batch["targets_segmentation"] != 0)
+    return jnp.sum(batch.get("targets_loss_mask", batch["targets_segmentation"]) != 0)
 
 
 class SFTDataHooks(BaseDataHooks):
diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py
index 9002caa16c..715585ec90 100644
--- a/src/maxtext/trainers/post_train/sft/train_sft.py
+++ b/src/maxtext/trainers/post_train/sft/train_sft.py
@@ -49,10 +49,12 @@
 
 from orbax import checkpoint as ocp
 
+from tunix.sft import diffusion as tunix_diffusion_sft
 from tunix.sft import metrics_logger, peft_trainer, profiler
 
 from maxtext.optimizers import optimizers
 from maxtext.configs import pyconfig
+from maxtext.integration.tunix import diffusion_sft as maxtext_diffusion_sft
 from maxtext.trainers.pre_train.train import loss_fn
 from maxtext.common.goodput import (
     GoodputEvent,
@@ -210,6 +212,9 @@ def loss_func(
       targets,
       targets_position,
       targets_segmentation,
+      completion_mask=None,
+      corruption_mask=None,
+      targets_loss_mask=None,
   ):
     data = {
         "inputs": inputs,
@@ -219,12 +224,36 @@ def loss_func(
         "targets_position": targets_position,
         "targets_segmentation": targets_segmentation,
     }
+    if completion_mask is not None:
+      data["completion_mask"] = completion_mask
+    if corruption_mask is not None:
+      data["corruption_mask"] = corruption_mask
+    if targets_loss_mask is not None:
+      data["targets_loss_mask"] = targets_loss_mask
     return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=True)
 
   trainer = trainer.with_loss_fn(loss_func, has_aux=True)
   return trainer
 
 
+def configure_training_objective(trainer, mt_config):
+  """Configures AR or target-aligned diffusion loss without changing AR defaults."""
+  if getattr(mt_config, "training_objective", "causal_lm") != "block_diffusion":
+    return use_maxtext_loss_function(trainer, mt_config)
+  max_logging.log("Configuring Tunix target-aligned block-diffusion SFT adapter.")
+  return tunix_diffusion_sft.configure_diffusion_sft(
+      trainer,
+      maxtext_diffusion_sft.create_batch_adapter(mt_config),
+      maxtext_diffusion_sft.create_target_aligned_logits_fn(mt_config),
+  )
+
+
+def _create_trainer(model, optimizer, tunix_config, mt_config):
+  if getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion":
+    return peft_trainer.PeftTrainer(model, optimizer, tunix_config)
+  return MaxTextPeftTrainer(model, optimizer, tunix_config)
+
+
 def validate_config(config):
   """Validates the configuration parameters for SFT training."""
   if config.optimizer_memory_host_offload:
@@ -237,6 +266,7 @@ def validate_config(config):
 
 def setup_trainer_state(mt_config, goodput_recorder=None):
   """Set up prerequisites for training loop."""
+  validate_config(mt_config)
   tunix_config = get_tunix_config(mt_config)
 
   with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT):
@@ -261,12 +291,12 @@ def setup_trainer_state(mt_config, goodput_recorder=None):
 
     # Provide rules context so 'norm' is translated to mesh axes during maybe_restore
     with nn_partitioning.axis_rules(mt_config.logical_axis_rules):
-      trainer = MaxTextPeftTrainer(model, optimizer, tunix_config)
+      trainer = _create_trainer(model, optimizer, tunix_config, mt_config)
       if mt_config.lora.lora_restore_path and trainer.train_steps == 0:
         lora_utils.restore_lora_from_path(trainer.model, mt_config)
       trainer.with_training_hooks(training_hooks)
       trainer.with_data_hooks(data_hooks)
-      trainer = use_maxtext_loss_function(trainer, mt_config)
+      trainer = configure_training_objective(trainer, mt_config)
 
   return trainer, mesh
 
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index ee92b93435..f27addd189 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -41,6 +41,7 @@
 from flax.nnx import variablelib
 
 from maxtext.configs import pyconfig
+from maxtext.diffusion import scoring as diffusion_scoring
 from maxtext.utils.globals import EPS
 from maxtext.utils import elastic_utils
 # Placeholder: internal
@@ -103,6 +104,22 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
     loss: average loss
     aux: a dictionary including intermediate_outputs, xent_sum, and total_weights
   """
+  is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion"
+  if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion:
+    raise ValueError(
+        "Block-diffusion attention requires target-aligned block-diffusion losses; "
+        "causal next-token labels would leak within a bidirectional block."
+    )
+  if is_block_diffusion:
+    required_masks = {"completion_mask", "corruption_mask", "targets_loss_mask"}
+    missing_masks = required_masks - data.keys()
+    if missing_masks:
+      raise ValueError(f"Block-diffusion loss requires explicit batch masks; missing {sorted(missing_masks)}")
+    target_shape = data["targets"].shape
+    for mask_name in required_masks:
+      if data[mask_name].shape != target_shape:
+        raise ValueError(f"{mask_name} must match targets shape; got {data[mask_name].shape} and {target_shape}")
+
   # decimate proportion of data when per_device_batch_size<1
   if is_train:
     for k, v in data.items():
@@ -110,6 +127,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
   else:
     for k, v in data.items():
       data[k] = v[: config.micro_batch_size_to_eval_on, :]
+  loss_mask_values = data["targets_loss_mask"] if is_block_diffusion else data["targets_segmentation"]
+  targets_loss_mask = (loss_mask_values != 0) & (data["targets_segmentation"] != 0)
+  target_positions = data.get("targets_position", data["inputs_position"])
+  if is_block_diffusion and config.block_diffusion_canvas_policy == "seed_and_mask":
+    if config.block_diffusion_logit_alignment == "shifted":
+      targets_loss_mask &= target_positions != 0
+    else:
+      targets_loss_mask &= target_positions % config.block_diffusion_block_size != 0
   mutable_collections = ["intermediates"]
   if config.mtp_num_layers > 0 and is_train:
     # The single model.apply call now triggers the entire chain if MTP is enabled:
@@ -161,6 +186,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
       hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
       xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train)
     else:
+      if is_block_diffusion:
+        logits = diffusion_scoring.align_logits_to_targets(
+            logits,
+            config.block_diffusion_logit_alignment,
+            target_positions,
+            data["targets_segmentation"] != 0,
+        )
       one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size)
       xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier)
 
@@ -179,9 +211,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
           debug_sharding=config.debug_sharding,
       )
 
-      # Mask out paddings at the end of each example.
-      xent = xent * (data["targets_segmentation"] != 0)
-      z_loss = z_loss * (data["targets_segmentation"] != 0)
+      xent = xent * targets_loss_mask
+      z_loss = z_loss * targets_loss_mask
 
       xent_sum = jnp.sum(xent)
       total_z_loss = jnp.sum(z_loss)
@@ -224,6 +255,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
       hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
       xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train)
     else:
+      if is_block_diffusion:
+        logits = diffusion_scoring.align_logits_to_targets(
+            logits,
+            config.block_diffusion_logit_alignment,
+            target_positions,
+            data["targets_segmentation"] != 0,
+        )
       one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size)
       xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier)
 
@@ -242,14 +280,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
           debug_sharding=config.debug_sharding,
       )
 
-      # Mask out paddings at the end of each example.
-      xent = xent * (data["targets_segmentation"] != 0)
-      z_loss = z_loss * (data["targets_segmentation"] != 0)
+      xent = xent * targets_loss_mask
+      z_loss = z_loss * targets_loss_mask
 
       xent_sum = jnp.sum(xent)
       total_z_loss = jnp.sum(z_loss)
 
-  total_weights = jnp.sum(data["targets_segmentation"] != 0)
+  total_weights = jnp.sum(targets_loss_mask)
   # If gradient accumulation is enabled, we don't need to divide xent_sum
   # by total_weights and then multiply the computed gradient by total_weights,
   # since it's equivalent to computing the gradient from xent_sum.
diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py
index 35fdf65503..eb0786c812 100644
--- a/src/maxtext/utils/gradient_accumulation.py
+++ b/src/maxtext/utils/gradient_accumulation.py
@@ -162,12 +162,15 @@ def reshape_to_microbatch_accumulations(batch_arr):
   grad_and_loss, aux = jax.lax.scan(
       accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps
   )
+  has_weights = grad_and_loss["total_weights"] > 0
+  safe_total_weights = jnp.maximum(grad_and_loss["total_weights"], 1)
   loss = (
-      grad_and_loss["loss"] / grad_and_loss["total_weights"]
+      grad_and_loss["loss"] / safe_total_weights
       + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps
       + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps
       + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps
   )
+  loss = jnp.where(has_weights, loss, 0.0)
   raw_grads = grad_and_loss["grad"]
   if data_parallel_active:
     # Mark the gradients unreduced over the "data" axis now that we're outside the
@@ -176,7 +179,10 @@ def reshape_to_microbatch_accumulations(batch_arr):
     unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings)
     raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings)
   raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings)
-  raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads)
+  raw_grads = jax.tree_util.tree_map(
+      lambda arr: jnp.where(has_weights, arr / safe_total_weights, jnp.zeros_like(arr)),
+      raw_grads,
+  )
   aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux)  # pytype: disable=module-attr
 
   if is_nnx:
diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py
index e6aebce059..943b1c63c2 100644
--- a/src/maxtext/utils/max_utils.py
+++ b/src/maxtext/utils/max_utils.py
@@ -957,6 +957,9 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu
       "targets_position",
       "inputs_segmentation",
       "targets_segmentation",
+      "completion_mask",
+      "corruption_mask",
+      "targets_loss_mask",
   }
 
   if hardware in ("gpu", "gpu_multiprocess"):
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 6c71e0487e..7f5845a72d 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -161,6 +161,10 @@ def get_shaped_batch(config, batch_sharding=None):
   shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
   shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
   shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
+  if getattr(config, "training_objective", "causal_lm") == "block_diffusion":
+    shaped_batch["completion_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
+    shaped_batch["corruption_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
+    shaped_batch["targets_loss_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding)
   if config.use_multimodal:
     image_shape = mm_processor.get_dummy_image_shape_for_init(
         config.model_name, batch_size=config.micro_batch_size_to_train_on
diff --git a/tests/post_training/unit/diffusion_sft_test.py b/tests/post_training/unit/diffusion_sft_test.py
new file mode 100644
index 0000000000..475b0ecdd7
--- /dev/null
+++ b/tests/post_training/unit/diffusion_sft_test.py
@@ -0,0 +1,95 @@
+# 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.
+
+"""Tests for the MaxText-to-Tunix diffusion SFT adapter."""
+
+from types import SimpleNamespace
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.diffusion import scoring
+from maxtext.integration.tunix import diffusion_sft
+
+
+pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only]
+
+
+def _config(alignment="shifted"):
+    return SimpleNamespace(
+        block_diffusion_block_size=4,
+        block_diffusion_logit_alignment=alignment,
+        enable_dropout=False,
+    )
+
+
+def _raw_batch():
+    positions = jnp.arange(8, dtype=jnp.int32)[None, :]
+    segmentation = jnp.ones((1, 8), dtype=jnp.int32)
+    return {
+        "inputs": jnp.asarray([[10, 11, 99, 99, 14, 99, 99, 99]], dtype=jnp.int32),
+        "inputs_position": positions,
+        "inputs_segmentation": segmentation,
+        "targets": jnp.asarray([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=jnp.int32),
+        "targets_position": positions,
+        "targets_segmentation": segmentation,
+        "completion_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32),
+        "corruption_mask": jnp.asarray([[0, 0, 1, 1, 0, 1, 1, 1]], dtype=jnp.int32),
+        "targets_loss_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32),
+    }
+
+
+def test_batch_adapter_preserves_shifted_anchor_weight():
+    batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch())
+
+    np.testing.assert_array_equal(batch.target_ids, _raw_batch()["targets"])
+    np.testing.assert_array_equal(batch.loss_weights, [[0, 0, 1, 1, 1, 1, 1, 1]])
+    assert batch.loss_weights.dtype == jnp.float32
+
+
+def test_batch_adapter_rejects_unowned_clean_target():
+    raw = _raw_batch()
+    raw["corruption_mask"] = raw["corruption_mask"].at[0, 3].set(0)
+
+    with pytest.raises(ValueError, match="corrupted targets or shifted block anchors"):
+        diffusion_sft.create_batch_adapter(_config())(raw)
+
+
+def test_logits_adapter_uses_target_alignment():
+    raw_logits = jnp.arange(1 * 8 * 3, dtype=jnp.float32).reshape(1, 8, 3)
+    calls = []
+
+    class Model:
+
+        def __call__(self, **kwargs):
+            calls.append(kwargs)
+            return raw_logits
+
+    batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch())
+    actual = diffusion_sft.create_target_aligned_logits_fn(_config())(
+        Model(), batch.model_inputs
+    )
+    expected = scoring.align_logits_to_targets(
+        raw_logits,
+        "shifted",
+        batch.model_inputs["target_positions"],
+        batch.model_inputs["target_segmentation"] != 0,
+    )
+
+    np.testing.assert_array_equal(actual, expected)
+    assert calls[0]["enable_dropout"] is False
+    np.testing.assert_array_equal(
+        calls[0]["decoder_input_tokens"], _raw_batch()["inputs"]
+    )
diff --git a/tests/post_training/unit/sft_hooks_test.py b/tests/post_training/unit/sft_hooks_test.py
index 9fde5add7e..8b97aceaa5 100644
--- a/tests/post_training/unit/sft_hooks_test.py
+++ b/tests/post_training/unit/sft_hooks_test.py
@@ -74,6 +74,18 @@ def test_sft_training_hooks_get_total_weights(self):
     total_weights = training_hooks.get_total_weights(batch)
     self.assertEqual(total_weights, 3)
 
+  def test_sft_training_hooks_prefers_explicit_loss_mask(self):
+    learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(self.config)
+    training_hooks = sft_hooks.SFTTrainingHooks(self.config, self.mesh, learning_rate_schedule, goodput_recorder=None)
+    batch = {
+        "targets_segmentation": np.ones((2, 4), dtype=np.int32),
+        "targets_loss_mask": np.array([[1, 0, 0, 0], [0, 1, 1, 0]], dtype=np.int32),
+    }
+
+    total_weights = training_hooks.get_total_weights(batch)
+
+    self.assertEqual(total_weights, 3)
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py
index f0df897b3e..f579f8159d 100644
--- a/tests/post_training/unit/train_sft_test.py
+++ b/tests/post_training/unit/train_sft_test.py
@@ -31,6 +31,7 @@ class TrainSFTTest(unittest.TestCase):
   def test_validate_config_valid(self):
     config = SimpleNamespace(
         optimizer_memory_host_offload=False,
+        eval_interval=-1,
     )
     # Should not raise any exception
     train_sft.validate_config(config)
@@ -43,6 +44,119 @@ def test_validate_config_invalid_offload(self):
     with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"):
       train_sft.validate_config(config)
 
+  @pytest.mark.cpu_only
+  def test_validate_config_accepts_weighted_diffusion_accumulation(self):
+    config = SimpleNamespace(
+        optimizer_memory_host_offload=False,
+        training_objective="block_diffusion",
+        gradient_accumulation_steps=2,
+        eval_interval=-1,
+    )
+
+    train_sft.validate_config(config)
+
+  @pytest.mark.cpu_only
+  def test_validate_config_accepts_weighted_diffusion_evaluation(self):
+    config = SimpleNamespace(
+        optimizer_memory_host_offload=False,
+        training_objective="block_diffusion",
+        gradient_accumulation_steps=1,
+        eval_interval=10,
+    )
+
+    train_sft.validate_config(config)
+
+  @pytest.mark.cpu_only
+  def test_setup_validates_before_initializing_model(self):
+    config = SimpleNamespace(
+        optimizer_memory_host_offload=True,
+    )
+
+    with (
+        mock.patch.object(train_sft, "get_tunix_config") as get_tunix_config,
+        mock.patch.object(train_sft.model_creation_utils, "from_pretrained") as from_pretrained,
+        self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True"),
+    ):
+      train_sft.setup_trainer_state(config)
+
+    get_tunix_config.assert_not_called()
+    from_pretrained.assert_not_called()
+
+  @pytest.mark.cpu_only
+  def test_maxtext_loss_wrapper_forwards_diffusion_masks(self):
+    trainer = mock.MagicMock()
+    trainer.with_loss_fn.return_value = trainer
+    config = SimpleNamespace()
+    train_sft.use_maxtext_loss_function(trainer, config)
+    wrapped_loss = trainer.with_loss_fn.call_args.args[0]
+    arrays = {
+        name: mock.sentinel.__getattr__(name)
+        for name in (
+            "inputs",
+            "inputs_position",
+            "inputs_segmentation",
+            "targets",
+            "targets_position",
+            "targets_segmentation",
+            "completion_mask",
+            "corruption_mask",
+            "targets_loss_mask",
+        )
+    }
+
+    with mock.patch.object(train_sft, "loss_fn", return_value=(1.0, {})) as loss_fn_mock:
+      wrapped_loss(mock.sentinel.model, **arrays)
+
+    forwarded = loss_fn_mock.call_args.args[2]
+    self.assertEqual(forwarded, arrays)
+    trainer.with_loss_fn.assert_called_once_with(wrapped_loss, has_aux=True)
+
+  @pytest.mark.cpu_only
+  def test_diffusion_objective_uses_tunix_adapter(self):
+    trainer = mock.sentinel.trainer
+    config = SimpleNamespace(training_objective="block_diffusion")
+    with (
+        mock.patch.object(train_sft.maxtext_diffusion_sft, "create_batch_adapter", return_value="adapter"),
+        mock.patch.object(
+            train_sft.maxtext_diffusion_sft,
+            "create_target_aligned_logits_fn",
+            return_value="logits_fn",
+        ),
+        mock.patch.object(
+            train_sft.tunix_diffusion_sft,
+            "configure_diffusion_sft",
+            return_value=trainer,
+        ) as configure,
+    ):
+      result = train_sft.configure_training_objective(trainer, config)
+
+    self.assertIs(result, trainer)
+    configure.assert_called_once_with(trainer, "adapter", "logits_fn")
+
+  @pytest.mark.cpu_only
+  def test_ar_objective_retains_maxtext_trainer_and_loss(self):
+    config = SimpleNamespace(training_objective="causal_lm")
+    with (
+        mock.patch.object(train_sft, "MaxTextPeftTrainer", return_value=mock.sentinel.trainer) as trainer_type,
+        mock.patch.object(train_sft, "use_maxtext_loss_function", return_value=mock.sentinel.configured) as configure,
+    ):
+      trainer = train_sft._create_trainer("model", "optimizer", "config", config)
+      result = train_sft.configure_training_objective(trainer, config)
+
+    self.assertIs(trainer, mock.sentinel.trainer)
+    self.assertIs(result, mock.sentinel.configured)
+    trainer_type.assert_called_once_with("model", "optimizer", "config")
+    configure.assert_called_once_with(mock.sentinel.trainer, config)
+
+  @pytest.mark.cpu_only
+  def test_diffusion_objective_uses_current_tunix_trainer(self):
+    config = SimpleNamespace(training_objective="block_diffusion")
+    with mock.patch.object(train_sft.peft_trainer, "PeftTrainer", return_value=mock.sentinel.trainer) as trainer_type:
+      trainer = train_sft._create_trainer("model", "optimizer", "config", config)
+
+    self.assertIs(trainer, mock.sentinel.trainer)
+    trainer_type.assert_called_once_with("model", "optimizer", "config")
+
   @pytest.mark.cpu_only
   def test_train_model_caching_moe(self):
     """Test that NNX graph caching is disabled for MoE models (num_experts > 1)."""
diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py
index 20522bbd3b..0c7aa03005 100644
--- a/tests/unit/attention_test.py
+++ b/tests/unit/attention_test.py
@@ -44,9 +44,12 @@
 from maxtext.layers import attention_op
 from maxtext.layers.attention_op import (
     AttentionOp,
+    BlockCausalMask,
     ChunkedCausalMask,
+    _generate_block_causal_attention_mask,
     _generate_chunk_attention_mask,
     _make_bidirectional_block_mask,
+    _resolve_attention_type,
 )
 from maxtext.layers.attentions import Attention
 from maxtext.layers import embeddings
@@ -330,6 +333,143 @@ def test_value_error_on_zero_chunk_size(self):
       _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0)
 
 
+class BlockCausalMaskTest(unittest.TestCase):
+  """Tests the shared dense and Splash block-causal masks."""
+
+  def _make_op(self, sequence_length, *, attention_type=AttentionType.BLOCK_DIFFUSION):
+    """Builds a minimal dot-product attention operator."""
+    config = types.SimpleNamespace(
+        block_diffusion_block_size=4,
+        context_parallel_load_balance=False,
+        context_sharding="context",
+    )
+    mesh = types.SimpleNamespace(shape={})
+    kwargs = {}
+    if attention_type is not None:
+      kwargs["attention_type"] = attention_type
+    return AttentionOp(
+        config=config,
+        num_query_heads=1,
+        num_kv_heads=1,
+        max_target_length=sequence_length,
+        mesh=mesh,
+        attention_kernel="dot_product",
+        **kwargs,
+    )
+
+  def test_dense_and_splash_masks_match(self):
+    sequence_length = 10
+    block_size = 4
+    query_positions = np.arange(sequence_length)[:, None]
+    key_positions = np.arange(sequence_length)[None, :]
+    expected = query_positions // block_size >= key_positions // block_size
+
+    splash_mask = BlockCausalMask((sequence_length, sequence_length), block_size)
+    dense_mask = _generate_block_causal_attention_mask((sequence_length, sequence_length), block_size)
+
+    np.testing.assert_array_equal(splash_mask[:, :], expected)
+    np.testing.assert_array_equal(dense_mask, expected)
+    self.assertTrue(expected[1, 3])
+    self.assertFalse(expected[3, 4])
+    self.assertTrue(expected[4, 3])
+    self.assertTrue(expected[8, 9])
+
+  def test_dot_product_mask_respects_packed_segments(self):
+    sequence_length = 12
+    query = jnp.zeros((1, sequence_length, 1, 8))
+    key = jnp.zeros((1, sequence_length, 1, 8))
+    segment_ids = jnp.asarray([[1] * 8 + [2] * 4], dtype=jnp.int32)
+
+    mask = self._make_op(sequence_length).generate_attention_mask(
+        query,
+        key,
+        segment_ids,
+        MODEL_MODE_TRAIN,
+    )
+
+    positions = np.arange(sequence_length)
+    expected = (positions[:, None] // 4 >= positions[None, :] // 4) & (
+        np.asarray(segment_ids[0])[:, None] == np.asarray(segment_ids[0])[None, :]
+    )
+    np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected)
+
+  def test_autoregressive_mask_is_unchanged(self):
+    key_length = 8
+    query = jnp.zeros((1, 1, 1, 8))
+    key = jnp.zeros((1, key_length, 1, 8))
+    segment_ids = jnp.asarray([[1, 1, 0, 1, 0, 0, 1, 0]], dtype=jnp.int32)
+
+    default_mask = self._make_op(key_length, attention_type=None).generate_attention_mask(
+        query,
+        key,
+        segment_ids,
+        MODEL_MODE_AUTOREGRESSIVE,
+    )
+    block_diffusion_mask = self._make_op(key_length).generate_attention_mask(
+        query,
+        key,
+        segment_ids,
+        MODEL_MODE_AUTOREGRESSIVE,
+    )
+
+    np.testing.assert_array_equal(block_diffusion_mask, default_mask)
+    np.testing.assert_array_equal(
+        np.asarray(default_mask == 0.0)[0, 0, 0, 0],
+        np.asarray(segment_ids[0] == DECODING_ACTIVE_SEQUENCE_INDICATOR),
+    )
+
+  def test_invalid_block_size_raises(self):
+    with self.assertRaises(ValueError):
+      BlockCausalMask((4, 4), 0)
+    with self.assertRaises(ValueError):
+      _generate_block_causal_attention_mask((4, 4), -1)
+
+
+class AttentionTypeResolutionTest(unittest.TestCase):
+
+  def test_config_selects_block_diffusion_without_model_dispatch(self):
+    config = types.SimpleNamespace(attention_type=AttentionType.BLOCK_DIFFUSION.value)
+
+    self.assertEqual(_resolve_attention_type(config, None), AttentionType.BLOCK_DIFFUSION)
+    self.assertEqual(_resolve_attention_type(config, AttentionType.GLOBAL), AttentionType.BLOCK_DIFFUSION)
+    self.assertEqual(_resolve_attention_type(config, AttentionType.FULL), AttentionType.FULL)
+    self.assertEqual(_resolve_attention_type(config, AttentionType.LOCAL_SLIDING), AttentionType.LOCAL_SLIDING)
+    self.assertEqual(
+        _resolve_attention_type(types.SimpleNamespace(attention_type=AttentionType.GLOBAL.value), AttentionType.FULL),
+        AttentionType.FULL,
+    )
+    self.assertEqual(_resolve_attention_type(types.SimpleNamespace(), None), AttentionType.GLOBAL)
+
+  def test_attention_op_honors_config_without_overriding_specialized_layers(self):
+    config = types.SimpleNamespace(
+        attention_type=AttentionType.BLOCK_DIFFUSION.value,
+        block_diffusion_block_size=4,
+    )
+    op = AttentionOp(
+        config=config,
+        num_query_heads=1,
+        num_kv_heads=1,
+        max_target_length=8,
+        mesh=types.SimpleNamespace(shape={}),
+        attention_kernel="dot_product",
+        attention_type=AttentionType.GLOBAL,
+    )
+
+    self.assertEqual(op.attention_type, AttentionType.BLOCK_DIFFUSION)
+
+    full_op = AttentionOp(
+        config=config,
+        num_query_heads=1,
+        num_kv_heads=1,
+        max_target_length=8,
+        mesh=types.SimpleNamespace(shape={}),
+        attention_kernel="dot_product",
+        attention_type=AttentionType.FULL,
+    )
+
+    self.assertEqual(full_op.attention_type, AttentionType.FULL)
+
+
 class LoadBalancedMaskTest(unittest.TestCase):
   """Tests for load-balanced Splash masks."""
 
@@ -370,6 +510,20 @@ def test_load_balanced_chunk_window(self):
 
     np.testing.assert_array_equal((causal_mask & chunk_mask)[:, :], expected_mask)
 
+  def test_load_balanced_block_causal_mask(self):
+    sequence_length = 8
+    block_size = 2
+    mask = attention_op.LoadBalancedBlockCausalMask(
+        shape=(sequence_length, sequence_length),
+        block_size=block_size,
+        cp_size=2,
+    )
+    query_positions = mask.q_sequence[:, None]
+    key_positions = np.arange(sequence_length)[None, :]
+    expected = query_positions // block_size >= key_positions // block_size
+
+    np.testing.assert_array_equal(mask[:, :], expected)
+
   def test_dot_product_local_mask_uses_segment_positions(self):
     config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context")
     mesh = types.SimpleNamespace(shape={"context": 4})
diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py
index f0504e4dd8..32a4bbba0e 100644
--- a/tests/unit/configs_value_test.py
+++ b/tests/unit/configs_value_test.py
@@ -192,6 +192,155 @@ def test_load_balanced_chunk_context_parallel_config(self):
     self.assertEqual(config.attention_type, "chunk")
     self.assertTrue(config.context_parallel_load_balance)
 
+  def test_block_diffusion_config(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "steps=1",
+        "attention=dot_product",
+        "attention_type=block_diffusion",
+        "block_diffusion_block_size=7",
+        "vocab_size=256",
+        "max_target_length=2048",
+        "hardware=cpu",
+    ]
+    config = pyconfig.initialize(argv)
+
+    self.assertEqual(config.attention_type, "block_diffusion")
+    self.assertEqual(config.block_diffusion_block_size, 7)
+    self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0)
+    self.assertFalse(hasattr(config, "enable_block_diffusion"))
+
+  def test_block_diffusion_invalid_configs_raise(self):
+    base_overrides = {
+        "run_name": "test",
+        "steps": 1,
+        "attention": "dot_product",
+        "attention_type": "block_diffusion",
+        "block_diffusion_block_size": 32,
+        "vocab_size": 256,
+        "max_target_length": 2048,
+        "hardware": "cpu",
+    }
+    cases = [
+        {"block_diffusion_block_size": 0},
+        {"attention": "cudnn_flash_te"},
+        {"attention": "flash", "hardware": "gpu"},
+    ]
+    for overrides in cases:
+      with self.subTest(overrides=overrides):
+        config = base_overrides | overrides
+        argv = ["", _BASE_CONFIG_PATH, *(f"{key}={value}" for key, value in config.items())]
+        with self.assertRaises((ValueError, pydantic.ValidationError)):
+          pyconfig.initialize(argv)
+
+  def test_default_attention_remains_global(self):
+    config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"])
+
+    self.assertEqual(config.attention_type, "global")
+    self.assertEqual(config.training_objective, "causal_lm")
+    self.assertEqual(config.block_diffusion_mask_id, -1)
+
+  def test_block_diffusion_training_objective_config(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "steps=1",
+        "training_objective=block_diffusion",
+        "attention=dot_product",
+        "attention_type=block_diffusion",
+        "block_diffusion_block_size=7",
+        "block_diffusion_mask_id=100",
+        "block_diffusion_min_noise=0.05",
+        "vocab_size=256",
+        "max_target_length=2048",
+        "packing=False",
+        "dataset_type=hf",
+        "hf_path=parquet",
+        "hardware=cpu",
+    ]
+
+    config = pyconfig.initialize(argv)
+
+    self.assertEqual(config.training_objective, "block_diffusion")
+    self.assertEqual(config.block_diffusion_mask_id, 100)
+    self.assertEqual(config.block_diffusion_min_noise, 0.05)
+    self.assertEqual(config.block_diffusion_logit_alignment, "same_position")
+    self.assertEqual(config.block_diffusion_canvas_policy, "all_masked")
+    self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0)
+
+  def test_block_diffusion_training_objective_rejects_incompatible_config(self):
+    base_overrides = {
+        "run_name": "test",
+        "steps": 1,
+        "training_objective": "block_diffusion",
+        "attention": "dot_product",
+        "attention_type": "block_diffusion",
+        "block_diffusion_block_size": 32,
+        "block_diffusion_mask_id": 100,
+        "block_diffusion_min_noise": 0.05,
+        "vocab_size": 256,
+        "max_target_length": 2048,
+        "packing": False,
+        "dataset_type": "hf",
+        "hf_path": "parquet",
+        "hardware": "cpu",
+    }
+    cases = [
+        ({"attention_type": "global"}, "attention_type='block_diffusion'"),
+        ({"block_diffusion_mask_id": -1}, "block_diffusion_mask_id"),
+        ({"block_diffusion_mask_id": 256}, "block_diffusion_mask_id"),
+        ({"block_diffusion_min_noise": 0.0}, "block_diffusion_min_noise"),
+        ({"packing": True}, "packing=False"),
+        ({"mtp_num_layers": 1}, "MTP"),
+        ({"num_vocab_tiling": 2}, "vocabulary tiling"),
+        ({"dataset_type": "grain"}, "dataset_type='hf'"),
+        ({"use_dpo": True}, "DPO"),
+        ({"use_multimodal": True}, "text-only"),
+        ({"block_diffusion_logit_alignment": "shifted"}, "seed_and_mask"),
+        (
+            {"block_diffusion_canvas_policy": "seed_and_mask"},
+            "same_position/all_masked",
+        ),
+    ]
+    for overrides, expected_regex in cases:
+      with self.subTest(overrides=overrides):
+        config = base_overrides | overrides
+        argv = [
+            "",
+            _BASE_CONFIG_PATH,
+            *(f"{key}={value}" for key, value in config.items()),
+        ]
+        with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex):
+          pyconfig.initialize(argv)
+
+  def test_shifted_block_diffusion_requires_seeded_canvas(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "steps=1",
+        "training_objective=block_diffusion",
+        "attention=dot_product",
+        "attention_type=block_diffusion",
+        "block_diffusion_block_size=8",
+        "block_diffusion_mask_id=100",
+        "block_diffusion_logit_alignment=shifted",
+        "block_diffusion_canvas_policy=seed_and_mask",
+        "vocab_size=256",
+        "packing=False",
+        "dataset_type=hf",
+        "hf_path=parquet",
+        "hardware=cpu",
+    ]
+
+    config = pyconfig.initialize(argv)
+
+    self.assertEqual(config.block_diffusion_logit_alignment, "shifted")
+    self.assertEqual(config.block_diffusion_canvas_policy, "seed_and_mask")
+
   @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"})
   def test_env_override(self):
     """Tests that environment variables override YAML values."""
diff --git a/tests/unit/diffusion_denoise_test.py b/tests/unit/diffusion_denoise_test.py
new file mode 100644
index 0000000000..37bf84491d
--- /dev/null
+++ b/tests/unit/diffusion_denoise_test.py
@@ -0,0 +1,267 @@
+# 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.
+
+"""Tests model-independent block-diffusion rollout transitions."""
+
+from absl.testing import absltest
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.diffusion import denoise
+
+
+def _target_logits(targets, vocab_size=32, high=12.0):
+  return jax.nn.one_hot(targets, vocab_size, dtype=jnp.float32) * high
+
+
+class DiffusionDenoiseTest(absltest.TestCase):
+
+  def test_same_position_generates_partial_blocks_and_preserves_prompt(self):
+    initial = jnp.asarray([[7, 8, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32)
+    expected = jnp.asarray([[7, 8, 12, 13, 14, 15, 16, 0]], dtype=jnp.int32)
+    positions = jnp.arange(8, dtype=jnp.int32)[None, :]
+    validity = initial != 0
+    completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_)
+
+    generated = denoise.low_confidence_generate(
+        lambda _: _target_logits(expected),
+        initial,
+        positions,
+        validity,
+        completion,
+        block_size=4,
+        mask_id=31,
+        logit_alignment="same_position",
+        canvas_policy="all_masked",
+    )
+
+    np.testing.assert_array_equal(generated, expected)
+
+  def test_shifted_seed_generates_anchors_before_each_block(self):
+    initial = jnp.asarray([[5, 6, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32)
+    expected = jnp.asarray([[5, 6, 9, 10, 11, 12, 13, 0]], dtype=jnp.int32)
+    positions = jnp.arange(8, dtype=jnp.int32)[None, :]
+    validity = initial != 0
+    completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_)
+
+    generated = denoise.low_confidence_generate(
+        lambda _: _target_logits(expected),
+        initial,
+        positions,
+        validity,
+        completion,
+        block_size=4,
+        mask_id=31,
+        logit_alignment="shifted",
+        canvas_policy="seed_and_mask",
+    )
+
+    np.testing.assert_array_equal(generated, expected)
+    self.assertEqual(int(generated[0, 4]), 11)
+
+  def test_forced_argmax_progress_completes_low_confidence_rows(self):
+    initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32)
+    positions = jnp.arange(4, dtype=jnp.int32)[None, :]
+    completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_)
+
+    generated = denoise.low_confidence_generate(
+        lambda canvas: jnp.zeros((*canvas.shape, 8), dtype=jnp.float32),
+        initial,
+        positions,
+        jnp.ones_like(completion),
+        completion,
+        block_size=4,
+        mask_id=7,
+        logit_alignment="same_position",
+        canvas_policy="all_masked",
+        confidence_threshold=0.99,
+    )
+
+    self.assertFalse(bool(jnp.any(generated == 7)))
+    np.testing.assert_array_equal(generated[:, :1], initial[:, :1])
+
+  def test_mask_token_is_never_committed(self):
+    initial = jnp.asarray([[6, 1, 1, 1]], dtype=jnp.int32)
+    positions = jnp.arange(4, dtype=jnp.int32)[None, :]
+    completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_)
+
+    def mask_favoring_logits(canvas):
+      logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32)
+      return logits.at[..., 7].set(100.0).at[..., 5].set(10.0)
+
+    generated = denoise.low_confidence_generate(
+        mask_favoring_logits,
+        initial,
+        positions,
+        jnp.ones_like(completion),
+        completion,
+        block_size=4,
+        mask_id=7,
+        logit_alignment="same_position",
+        canvas_policy="all_masked",
+    )
+
+    np.testing.assert_array_equal(generated, [[6, 5, 5, 5]])
+
+  def test_mask_token_is_excluded_from_shifted_anchor_proposals(self):
+    initial = jnp.asarray([[2, 3, 4, 5, 1, 1]], dtype=jnp.int32)
+    positions = jnp.arange(6, dtype=jnp.int32)[None, :]
+    completion = jnp.asarray([[0, 0, 0, 0, 1, 1]], dtype=jnp.bool_)
+
+    def mask_favoring_logits(canvas):
+      logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32)
+      return logits.at[..., 7].set(100.0).at[..., 6].set(10.0)
+
+    generated = denoise.low_confidence_generate(
+        mask_favoring_logits,
+        initial,
+        positions,
+        jnp.ones_like(completion),
+        completion,
+        block_size=4,
+        mask_id=7,
+        logit_alignment="shifted",
+        canvas_policy="seed_and_mask",
+    )
+
+    np.testing.assert_array_equal(generated, [[2, 3, 4, 5, 6, 6]])
+
+  def test_rollout_is_jittable(self):
+    positions = jnp.arange(4, dtype=jnp.int32)[None, :]
+    completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_)
+    expected = jnp.asarray([[4, 5, 6, 7]], dtype=jnp.int32)
+
+    generate = jax.jit(
+        lambda initial: denoise.low_confidence_generate(
+            lambda _: _target_logits(expected),
+            initial,
+            positions,
+            jnp.ones_like(completion),
+            completion,
+            block_size=4,
+            mask_id=31,
+            logit_alignment="same_position",
+            canvas_policy="all_masked",
+        )
+    )
+
+    np.testing.assert_array_equal(generate(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32)), expected)
+
+  def test_logical_positions_drive_blocks_after_reordering(self):
+    positions = jnp.asarray([[0, 4, 1, 5, 2, 6, 3, 7]], dtype=jnp.int32)
+    expected = positions + 10
+    initial = jnp.where(positions < 2, expected, 1)
+    completion = positions >= 2
+
+    generated = denoise.low_confidence_generate(
+        lambda _: _target_logits(expected),
+        initial,
+        positions,
+        jnp.ones_like(completion),
+        completion,
+        block_size=4,
+        mask_id=31,
+        logit_alignment="same_position",
+        canvas_policy="all_masked",
+    )
+
+    np.testing.assert_array_equal(generated, expected)
+
+  def test_forced_ties_are_invariant_to_physical_reordering(self):
+    completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_)
+
+    def generate(physical_order):
+      positions = jnp.asarray(physical_order, dtype=jnp.int32)[None, :]
+      initial = jnp.take_along_axis(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32), positions, axis=1)
+      physical_completion = jnp.take_along_axis(completion, positions, axis=1)
+
+      def low_confidence_logits(canvas):
+        committed_count = jnp.sum((canvas != 31) & physical_completion, axis=1)
+        proposed_token = jnp.broadcast_to((10 + committed_count)[:, None], canvas.shape)
+        return _target_logits(proposed_token, vocab_size=32, high=0.1)
+
+      output = denoise.low_confidence_generate(
+          low_confidence_logits,
+          initial,
+          positions,
+          jnp.ones_like(physical_completion),
+          physical_completion,
+          block_size=4,
+          mask_id=31,
+          logit_alignment="same_position",
+          canvas_policy="all_masked",
+          confidence_threshold=0.99,
+      )
+      inverse_order = jnp.argsort(positions, axis=1)
+      return jnp.take_along_axis(output, inverse_order, axis=1)
+
+    np.testing.assert_array_equal(generate([0, 1, 2, 3]), generate([0, 3, 1, 2]))
+
+  def test_rejects_non_suffix_and_shifted_origin_completion(self):
+    positions = np.arange(4, dtype=np.int32)[None, :]
+    validity = np.ones((1, 4), dtype=bool)
+
+    with self.assertRaisesRegex(ValueError, "contiguous suffix"):
+      denoise.validate_completion_suffix(positions, validity, np.asarray([[0, 1, 0, 1]], dtype=bool))
+    with self.assertRaisesRegex(ValueError, "position zero"):
+      denoise.validate_completion_suffix(
+          positions,
+          validity,
+          np.asarray([[1, 1, 1, 1]], dtype=bool),
+          shifted_seed=True,
+      )
+
+  def test_rejects_missing_duplicate_or_out_of_range_positions(self):
+    validity = np.ones((1, 3), dtype=bool)
+    completion = np.asarray([[0, 1, 1]], dtype=bool)
+
+    for positions in (
+        np.asarray([[1, 2, 3]], dtype=np.int32),
+        np.asarray([[0, 1, 1]], dtype=np.int32),
+        np.asarray([[0, 1, 5]], dtype=np.int32),
+    ):
+      with self.subTest(positions=positions), self.assertRaisesRegex(ValueError, "logical positions"):
+        denoise.validate_completion_suffix(positions, validity, completion)
+
+  def test_rejects_unsupported_contract_and_short_step_cap(self):
+    tokens = jnp.ones((1, 4), dtype=jnp.int32)
+    positions = jnp.arange(4, dtype=jnp.int32)[None, :]
+    mask = jnp.ones_like(tokens, dtype=jnp.bool_)
+
+    with self.assertRaisesRegex(ValueError, "supports only"):
+      denoise.low_confidence_generate(
+          lambda _: _target_logits(tokens),
+          tokens,
+          positions,
+          mask,
+          mask,
+          block_size=4,
+          mask_id=31,
+          logit_alignment="same_position",
+          canvas_policy="seed_and_mask",
+      )
+    with self.assertRaisesRegex(ValueError, "at least block_size"):
+      denoise.low_confidence_generate(
+          lambda _: _target_logits(tokens),
+          tokens,
+          positions,
+          mask,
+          mask,
+          block_size=4,
+          mask_id=31,
+          logit_alignment="same_position",
+          canvas_policy="all_masked",
+          max_denoise_steps=3,
+      )
diff --git a/tests/unit/gradient_accumulation_nnx_test.py b/tests/unit/gradient_accumulation_nnx_test.py
index 478fe9a746..5ca56f484d 100644
--- a/tests/unit/gradient_accumulation_nnx_test.py
+++ b/tests/unit/gradient_accumulation_nnx_test.py
@@ -14,6 +14,8 @@
 
 """Unit tests for the NNX branch of gradient_accumulation_loss_and_grad."""
 
+# pylint: disable=too-many-positional-arguments
+
 import unittest
 from dataclasses import dataclass
 
@@ -67,6 +69,20 @@ def _fake_loss_fn(model, config, data, dropout_rng, params, is_train=True):
   return xent_sum / total_weights, aux
 
 
+def _zero_weight_loss_fn(model, config, data, dropout_rng, params, is_train=True):
+  """Produces model-dependent sums with a zero token denominator."""
+  del config, dropout_rng, params, is_train
+  xent_sum = jnp.sum(model(data["inputs"]) ** 2) + 1.0
+  aux = {
+      "xent_sum": xent_sum,
+      "total_weights": jnp.array(0.0),
+      "moe_lb_loss": jnp.array(0.0),
+      "indexer_loss": jnp.array(0.0),
+      "mtp_loss": jnp.array(0.0),
+  }
+  return xent_sum, aux
+
+
 class TestGradientAccumulationNNX(unittest.TestCase):
   """Cover the NNX path of gradient_accumulation_loss_and_grad."""
 
@@ -151,6 +167,22 @@ def test_nnx_with_shard_optimizer_over_data_casts_to_bf16(self):
     for g in jax.tree.leaves(raw_grads):
       self.assertTrue(jnp.all(jnp.isfinite(g)))
 
+  def test_zero_total_weights_returns_zero_loss_and_gradients(self):
+    loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad(
+        _zero_weight_loss_fn,
+        self.cfg,
+        self.model,
+        params=None,
+        params_shardings=self._params_shardings(),
+        data=self.data,
+        dropout_rng=None,
+    )
+
+    self.assertEqual(float(loss), 0.0)
+    for gradient in jax.tree.leaves(raw_grads):
+      self.assertTrue(jnp.all(jnp.isfinite(gradient)))
+      self.assertTrue(jnp.all(gradient == 0))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py
index 93564514a3..a23525af02 100644
--- a/tests/unit/hf_data_processing_test.py
+++ b/tests/unit/hf_data_processing_test.py
@@ -15,6 +15,7 @@
 """Tests for Hugging Face data processing."""
 
 import sys
+from types import SimpleNamespace
 import unittest
 import os.path
 
@@ -26,6 +27,7 @@
 from maxtext.configs import pyconfig
 from maxtext.input_pipeline import hf_data_processing
 from maxtext.input_pipeline import input_pipeline_interface
+from maxtext.input_pipeline import input_pipeline_utils
 from maxtext.common.gcloud_stub import is_decoupled
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
 from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory
@@ -126,5 +128,75 @@ def get_first_batch(iterator):
     self.assertTrue((train_batch1["targets"] == train_batch2["targets"]).all())  # pytype: disable=unsupported-operands
 
 
+@pytest.mark.cpu_only
+class TrainingObjectiveTransformTest(unittest.TestCase):
+  """Tests objective selection at the HF pipeline boundary."""
+
+  def test_default_objective_keeps_shift_data(self):
+    transform = hf_data_processing._get_training_objective_transform(  # pylint: disable=protected-access
+        SimpleNamespace(),
+        shift=True,
+        use_dpo=False,
+        use_sft=False,
+        completion_only=False,
+        packing=False,
+        pad_id=0,
+        bos_token_id=1,
+    )
+
+    self.assertIsInstance(transform, input_pipeline_utils.ShiftData)
+    self.assertEqual(transform.ignored_ids, [0, 1])
+
+  def test_explicit_block_diffusion_objective_replaces_shift(self):
+    config = SimpleNamespace(
+        training_objective="block_diffusion",
+        block_diffusion_block_size=4,
+        block_diffusion_mask_id=99,
+        block_diffusion_min_noise=0.05,
+        block_diffusion_canvas_policy="seed_and_mask",
+        block_diffusion_logit_alignment="shifted",
+    )
+
+    transform = hf_data_processing._get_training_objective_transform(  # pylint: disable=protected-access
+        config,
+        shift=True,
+        use_dpo=False,
+        use_sft=True,
+        completion_only=True,
+        packing=False,
+        pad_id=0,
+        bos_token_id=1,
+    )
+
+    self.assertIsInstance(transform, input_pipeline_utils.BlockDiffusionCorruption)
+    self.assertEqual(transform.block_size, 4)
+    self.assertEqual(transform.mask_id, 99)
+    self.assertEqual(transform.min_noise, 0.05)
+    self.assertTrue(transform.completion_only)
+    self.assertTrue(transform.seed_first_token)
+    self.assertTrue(transform.include_seed_in_loss)
+
+  def test_block_diffusion_rejects_packing(self):
+    config = SimpleNamespace(
+        training_objective="block_diffusion",
+        block_diffusion_block_size=4,
+        block_diffusion_mask_id=99,
+        block_diffusion_min_noise=0.05,
+        block_diffusion_canvas_policy="all_masked",
+    )
+
+    with self.assertRaisesRegex(ValueError, "packing=False"):
+      hf_data_processing._get_training_objective_transform(  # pylint: disable=protected-access
+          config,
+          shift=True,
+          use_dpo=False,
+          use_sft=False,
+          completion_only=False,
+          packing=True,
+          pad_id=0,
+          bos_token_id=1,
+      )
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py
index 0d5adae4a9..fe5ce30cda 100644
--- a/tests/unit/input_pipeline_utils_test.py
+++ b/tests/unit/input_pipeline_utils_test.py
@@ -17,7 +17,14 @@
 import pytest
 import unittest
 
-from maxtext.input_pipeline.input_pipeline_utils import compute_file_sharding
+import numpy as np
+
+from maxtext.input_pipeline.input_pipeline_utils import (
+    BlockDiffusionCorruption,
+    PadOrTrimToMaxLength,
+    SFTPromptMasking,
+    compute_file_sharding,
+)
 
 
 @pytest.mark.cpu_only
@@ -109,5 +116,210 @@ def test_no_row_shard_when_only_one_reader(self):
     self.assertIsNone(row_shard)
 
 
+def _make_block_diffusion_batch():
+  """Builds a padded, completion-only SFT batch with a partial final block."""
+  inputs = np.asarray(
+      [
+          [11, 12, 21, 22, 23, 24, 25, 26, 27, 0],
+          [31, 41, 42, 43, 44, 45, 0, 0, 0, 0],
+      ],
+      dtype=np.int32,
+  )
+  targets = np.asarray(
+      [
+          [11, 12, 21, 22, 23, 24, 25, 26, 27, 0],
+          [31, 41, 42, 43, 44, 45, 0, 0, 0, 0],
+      ],
+      dtype=np.int32,
+  )
+  completion_mask = np.asarray(
+      [
+          [0, 0, 1, 1, 1, 1, 1, 1, 1, 0],
+          [0, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+      ],
+      dtype=np.int32,
+  )
+  positions = np.broadcast_to(np.arange(inputs.shape[1], dtype=np.int32), inputs.shape).copy()
+  return {
+      "inputs": inputs,
+      "targets": targets,
+      "inputs_segmentation": (inputs != 0).astype(np.int32),
+      "targets_segmentation": (targets != 0).astype(np.int32),
+      "inputs_position": positions.copy(),
+      "targets_position": positions.copy(),
+      "completion_mask": completion_mask,
+  }
+
+
+class _ControlledRng:
+  """Supplies distinct noise rates while making token draws deterministic."""
+
+  def __init__(self):
+    self.uniform_calls = 0
+    self.random_calls = 0
+
+  def uniform(self, low, high, size):
+    del low, high
+    self.uniform_calls += 1
+    return np.asarray([[[0.1], [0.9], [0.1]], [[0.1], [0.9], [0.1]]]).reshape(size)
+
+  def random(self, size):
+    self.random_calls += 1
+    return np.full(size, 0.5)
+
+
+@pytest.mark.cpu_only
+class BlockDiffusionCorruptionTest(unittest.TestCase):
+  """Tests aligned block-diffusion corruption metadata."""
+
+  def _apply(self, seed=0):
+    transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0e-3, completion_only=True)
+    return transform.random_map(_make_block_diffusion_batch(), np.random.default_rng(seed))
+
+  def test_preserves_targets_and_segmentation(self):
+    clean = _make_block_diffusion_batch()
+    output = self._apply()
+
+    for key in (
+        "targets",
+        "inputs_segmentation",
+        "targets_segmentation",
+        "inputs_position",
+        "targets_position",
+    ):
+      np.testing.assert_array_equal(output[key], clean[key])
+    np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"])
+    self.assertTrue(np.all(output["targets_loss_mask"] <= output["completion_mask"]))
+    self.assertFalse(np.shares_memory(output["completion_mask"], output["targets_loss_mask"]))
+    np.testing.assert_array_equal(output["corruption_mask"], output["targets_loss_mask"])
+
+    loss_positions = output["targets_loss_mask"] != 0
+    self.assertTrue(np.all(output["inputs"][loss_positions] == 99))
+    np.testing.assert_array_equal(output["inputs"][~loss_positions], clean["inputs"][~loss_positions])
+
+  def test_every_eligible_full_or_partial_block_has_loss(self):
+    output = self._apply(seed=17)
+    completion_mask = output["completion_mask"] != 0
+    loss_mask = output["targets_loss_mask"] != 0
+
+    for row in range(completion_mask.shape[0]):
+      for start in range(0, completion_mask.shape[1], 4):
+        stop = min(start + 4, completion_mask.shape[1])
+        if np.any(completion_mask[row, start:stop]):
+          self.assertTrue(np.any(loss_mask[row, start:stop]), (row, start, stop))
+    self.assertTrue(loss_mask[0, 8])
+
+  def test_corruption_rate_is_sampled_per_block(self):
+    rng = _ControlledRng()
+    transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True)
+
+    output = transform.random_map(_make_block_diffusion_batch(), rng)
+
+    self.assertEqual(rng.uniform_calls, 1)
+    self.assertEqual(rng.random_calls, 2)
+    self.assertEqual(int(output["targets_loss_mask"][0, :4].sum()), 1)
+    self.assertEqual(int(output["targets_loss_mask"][0, 4:8].sum()), 4)
+
+  def test_deterministic_for_fixed_seed(self):
+    output_a = self._apply(seed=123)
+    output_b = self._apply(seed=123)
+    for key, value in output_a.items():
+      np.testing.assert_array_equal(value, output_b[key])
+
+  def test_invalid_parameters_raise(self):
+    with self.assertRaises(ValueError):
+      BlockDiffusionCorruption(block_size=0, mask_id=99)
+    with self.assertRaises(ValueError):
+      BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=0.0)
+    with self.assertRaisesRegex(ValueError, "requires seed_first_token"):
+      BlockDiffusionCorruption(block_size=4, mask_id=99, include_seed_in_loss=True)
+
+  def test_completion_only_requires_explicit_role_mask(self):
+    clean = _make_block_diffusion_batch()
+    del clean["completion_mask"]
+
+    with self.assertRaisesRegex(ValueError, "explicit completion_mask"):
+      self._apply_transform(clean)
+
+  def test_completion_mask_does_not_replace_full_validity(self):
+    clean = _make_block_diffusion_batch()
+    transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0, completion_only=False)
+
+    output = transform.random_map(clean, np.random.default_rng(0))
+
+    np.testing.assert_array_equal(output["targets_segmentation"], clean["targets_segmentation"])
+    np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"])
+    self.assertTrue(np.all(output["targets_loss_mask"][clean["inputs_segmentation"] != 0]))
+    self.assertTrue(np.any(output["targets_loss_mask"] > output["completion_mask"]))
+
+  def test_all_padding_has_finite_zero_masks(self):
+    clean = _make_block_diffusion_batch()
+    for key in (
+        "inputs",
+        "targets",
+        "inputs_segmentation",
+        "targets_segmentation",
+        "completion_mask",
+    ):
+      clean[key] = np.zeros_like(clean[key])
+
+    output = self._apply_transform(clean)
+
+    self.assertFalse(output["corruption_mask"].any())
+    self.assertFalse(output["targets_loss_mask"].any())
+
+  def test_shifted_seed_canvas_keeps_anchors_clean_but_supervised(self):
+    clean = _make_block_diffusion_batch()
+    transform = BlockDiffusionCorruption(
+        block_size=4,
+        mask_id=99,
+        min_noise=1.0,
+        completion_only=False,
+        seed_first_token=True,
+        include_seed_in_loss=True,
+    )
+
+    output = transform.random_map(clean, np.random.default_rng(0))
+
+    self.assertFalse(output["corruption_mask"][:, ::4].any())
+    expected_loss = clean["targets_segmentation"].copy()
+    expected_loss[:, 0] = 0
+    np.testing.assert_array_equal(output["targets_loss_mask"], expected_loss)
+    self.assertTrue(output["targets_loss_mask"][0, 4])
+    self.assertTrue(output["targets_loss_mask"][0, 8])
+
+  def _apply_transform(self, batch):
+    transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True)
+    return transform.random_map(batch, np.random.default_rng(0))
+
+
+@pytest.mark.cpu_only
+class BlockDiffusionSftMaskTest(unittest.TestCase):
+  """Tests role metadata before randomized corruption."""
+
+  def test_target_aligned_sft_preserves_clean_tokens_and_completion_roles(self):
+    transform = SFTPromptMasking(
+        text_column_name="messages",
+        completion_only=True,
+        max_target_length=8,
+        unk_id=0,
+        target_aligned=True,
+        emit_completion_mask=True,
+    )
+    element = {
+        "messages": [[11, 12], [21, 22], [31], [41, 42]],
+        "is_prompt": [True, False, True, False],
+    }
+
+    output = transform.map(element)
+    padded = PadOrTrimToMaxLength(max_length=8, pad_id=0).map(output)
+
+    np.testing.assert_array_equal(padded["targets"], padded["inputs"])
+    np.testing.assert_array_equal(padded["completion_mask"], [0, 0, 1, 1, 0, 1, 1, 0])
+    np.testing.assert_array_equal(padded["targets_segmentation"], [1, 1, 1, 1, 1, 1, 1, 0])
+    self.assertNotIn("completion_mask_segmentation", padded)
+    self.assertNotIn("completion_mask_position", padded)
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py
index 097d85d491..7dcc6cea83 100644
--- a/tests/unit/max_utils_test.py
+++ b/tests/unit/max_utils_test.py
@@ -23,10 +23,12 @@
 import jax
 from jax import numpy as jnp
 from jax import random
+from maxtext.common.common_types import ReorderStrategy
 from maxtext.configs import pyconfig
 from maxtext.utils import max_utils
 from maxtext.utils.train_utils import setup_train_loop
 from tests.utils.test_helpers import get_test_config_path
+import numpy as np
 import optax
 import pytest
 
@@ -655,6 +657,31 @@ def test_reorder_roundtrip(self):
     # 3. Assert roundtrip is lossless
     self.assertTrue(jnp.allclose(x, restored, rtol=1e-5, atol=1e-6))
 
+  def test_block_diffusion_masks_are_reordered_with_tokens(self):
+    completion_mask = jnp.arange(16, dtype=jnp.int32).reshape(1, 16)
+    targets_loss_mask = completion_mask + 100
+    corruption_mask = completion_mask + 50
+    batch = {
+        "inputs": completion_mask + 200,
+        "completion_mask": completion_mask,
+        "corruption_mask": corruption_mask,
+        "targets_loss_mask": targets_loss_mask,
+    }
+
+    reordered = max_utils.reorder_causal_load_balanced(
+        batch,
+        cp_size=2,
+        reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP,
+        hardware="cpu",
+    )
+
+    np.testing.assert_array_equal(reordered["completion_mask"], max_utils.reorder_sequence(completion_mask, 2))
+    np.testing.assert_array_equal(reordered["corruption_mask"], max_utils.reorder_sequence(corruption_mask, 2))
+    np.testing.assert_array_equal(
+        reordered["targets_loss_mask"],
+        max_utils.reorder_sequence(targets_loss_mask, 2),
+    )
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py
index 28c9f6572a..7f33735c1e 100644
--- a/tests/unit/maxtext_utils_test.py
+++ b/tests/unit/maxtext_utils_test.py
@@ -1140,13 +1140,22 @@ def test_linen_in_shardings_includes_rng(self):
 class TestGetShapedBatch(unittest.TestCase):
   """Tests for get_shaped_batch."""
 
-  def _make_cfg(self, *, enable_diloco=False, use_multimodal=False, use_audio=False):
+  def _make_cfg(
+      self,
+      *,
+      enable_diloco=False,
+      use_multimodal=False,
+      use_audio=False,
+      training_objective="causal_lm",
+  ):
+    """Builds the config subset consumed by get_shaped_batch."""
     cfg = MagicMock()
     cfg.enable_diloco = enable_diloco
     cfg.global_batch_size_to_load = 4
     cfg.max_target_length = 16
     cfg.use_multimodal = use_multimodal
     cfg.use_audio = use_audio
+    cfg.training_objective = training_objective
     if enable_diloco:
       cfg.num_diloco_replicas = 2
     return cfg
@@ -1169,6 +1178,23 @@ def test_standard_shape(self):
     expected_shape = (cfg.global_batch_size_to_load, cfg.max_target_length)
     self.assertEqual(batch["inputs"].shape, expected_shape)
 
+  def test_block_diffusion_masks_are_in_shaped_batch(self):
+    cfg = self._make_cfg(training_objective="block_diffusion")
+
+    batch = maxtext_utils.get_shaped_batch(cfg)
+
+    self.assertEqual(batch["completion_mask"].shape, batch["inputs"].shape)
+    self.assertEqual(batch["corruption_mask"].shape, batch["inputs"].shape)
+    self.assertEqual(batch["targets_loss_mask"].shape, batch["inputs"].shape)
+    self.assertEqual(batch["targets_loss_mask"].dtype, jnp.int32)
+
+  def test_causal_shaped_batch_has_no_diffusion_masks(self):
+    batch = maxtext_utils.get_shaped_batch(self._make_cfg())
+
+    self.assertNotIn("completion_mask", batch)
+    self.assertNotIn("corruption_mask", batch)
+    self.assertNotIn("targets_loss_mask", batch)
+
   def test_diloco_shape(self):
     cfg = self._make_cfg(enable_diloco=True)
     batch = maxtext_utils.get_shaped_batch(cfg)
diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py
new file mode 100644
index 0000000000..07e4f6200b
--- /dev/null
+++ b/tests/unit/pre_train_loss_mask_test.py
@@ -0,0 +1,264 @@
+# 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.
+
+"""Tests explicit target-loss masking in the pre-training loss."""
+
+# pylint: disable=too-many-positional-arguments
+
+from dataclasses import dataclass
+import unittest
+from unittest import mock
+
+from flax import linen as nn
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.trainers.pre_train import train as pre_train
+from maxtext.diffusion import scoring as diffusion_scoring
+
+
+@dataclass
+class _Config:
+  """Configuration subset consumed by pre-training loss_fn."""
+
+  micro_batch_size_to_train_on: int = 2
+  micro_batch_size_to_eval_on: int = 2
+  vocab_size: int = 8
+  z_loss_multiplier: float = 0.0
+  enable_dropout: bool = False
+  use_multimodal: bool = False
+  use_indexer: bool = False
+  indexer_sparse_training: bool = False
+  indexer_loss_scaling_factor: float = 0.0
+  num_vocab_tiling: int = 1
+  num_experts: int = 1
+  routed_bias: bool = False
+  routed_bias_update_rate: float = 0.0
+  mtp_num_layers: int = 0
+  mtp_eval_target_module: int = 0
+  use_qk_clip: bool = False
+  use_tunix_gradient_accumulation: bool = False
+  gradient_accumulation_steps: int = 1
+  shard_mode: int = 0
+  debug_sharding: bool = False
+  weight_sparsity_n: int = 0
+  weight_sparsity_m: int = 0
+  attention_type: str = "global"
+  training_objective: str = "causal_lm"
+  block_diffusion_logit_alignment: str = "same_position"
+  block_diffusion_canvas_policy: str = "all_masked"
+  block_diffusion_block_size: int = 4
+
+
+class _UniformNnxDecoder(nnx.Module):
+  """Returns uniform logits through the NNX call contract."""
+
+  def __init__(self, vocab_size):
+    self.vocab_size = vocab_size
+    self.mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context"))
+
+  def __call__(
+      self,
+      decoder_input_tokens,
+      decoder_positions,
+      decoder_segment_ids=None,
+      encoder_images=None,
+      encoder_image_masks=None,
+      enable_dropout=False,
+      decoder_target_tokens=None,
+      decoder_target_mask=None,
+  ):
+    del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks
+    del enable_dropout, decoder_target_tokens, decoder_target_mask
+    return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32)
+
+
+class _UniformLinenDecoder(nn.Module):
+  """Returns uniform logits through the Linen call contract."""
+
+  vocab_size: int
+  mesh: object
+
+  @nn.compact
+  def __call__(
+      self,
+      decoder_input_tokens,
+      decoder_positions,
+      decoder_segment_ids=None,
+      encoder_images=None,
+      encoder_image_masks=None,
+      enable_dropout=False,
+      decoder_target_tokens=None,
+      decoder_target_mask=None,
+  ):
+    del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks
+    del enable_dropout, decoder_target_tokens, decoder_target_mask
+    return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32)
+
+
+def _make_data(include_loss_mask=True):
+  """Builds a batch whose explicit loss mask differs from segmentation."""
+  data = {
+      "inputs": jnp.zeros((2, 4), dtype=jnp.int32),
+      "inputs_position": jnp.broadcast_to(jnp.arange(4), (2, 4)),
+      "inputs_segmentation": jnp.ones((2, 4), dtype=jnp.int32),
+      "targets": jnp.zeros((2, 4), dtype=jnp.int32),
+      "targets_segmentation": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32),
+      "completion_mask": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32),
+      "corruption_mask": jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32),
+  }
+  if include_loss_mask:
+    data["targets_loss_mask"] = jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32)
+  return data
+
+
+class PreTrainLossMaskTest(unittest.TestCase):
+  """Checks explicit masks in both Linen and NNX loss branches."""
+
+  def setUp(self):
+    super().setUp()
+    self.config = _Config()
+    self.per_token_xent = jnp.arange(1, 9, dtype=jnp.float32).reshape(2, 4)
+    self.per_token_z_loss = self.per_token_xent / 10.0
+
+  def _cross_entropy_patch(self):
+    return mock.patch.object(
+        pre_train.max_utils,
+        "cross_entropy_with_logits",
+        return_value=(self.per_token_xent, self.per_token_z_loss),
+    )
+
+  def _use_block_diffusion(self):
+    self.config.attention_type = "block_diffusion"
+    self.config.training_objective = "block_diffusion"
+
+  def _assert_explicit_mask_result(self, loss, aux):
+    expected_mask = _make_data()["targets_loss_mask"] != 0
+    expected_xent = jnp.sum(self.per_token_xent * expected_mask)
+    expected_z_loss = jnp.sum(self.per_token_z_loss * expected_mask) / 3.0
+    self.assertEqual(int(aux["total_weights"]), 3)
+    self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent))
+    self.assertAlmostEqual(float(loss), float(expected_xent / 3.0))
+    self.assertAlmostEqual(float(aux["z_loss"]), float(expected_z_loss))
+
+  def test_nnx_loss_uses_targets_loss_mask(self):
+    self._use_block_diffusion()
+    model = _UniformNnxDecoder(self.config.vocab_size)
+    with self._cross_entropy_patch():
+      loss, aux = pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True)
+
+    self._assert_explicit_mask_result(loss, aux)
+
+  def test_linen_loss_uses_targets_loss_mask(self):
+    self._use_block_diffusion()
+    mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context"))
+    model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh)
+    data = _make_data()
+    variables = model.init(
+        jax.random.key(0),
+        data["inputs"],
+        data["inputs_position"],
+        decoder_segment_ids=data["inputs_segmentation"],
+        decoder_target_tokens=data["targets"],
+        decoder_target_mask=data["targets_segmentation"],
+    )
+    with self._cross_entropy_patch():
+      loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True)
+
+    self._assert_explicit_mask_result(loss, aux)
+
+  def test_causal_fallback_uses_targets_segmentation(self):
+    data = _make_data(include_loss_mask=False)
+    model = _UniformNnxDecoder(self.config.vocab_size)
+    with self._cross_entropy_patch():
+      loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True)
+
+    expected_mask = data["targets_segmentation"] != 0
+    expected_xent = jnp.sum(self.per_token_xent * expected_mask)
+    self.assertEqual(int(aux["total_weights"]), 7)
+    self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent))
+    self.assertAlmostEqual(float(loss), float(expected_xent / 7.0))
+
+  def test_zero_explicit_mask_has_finite_zero_loss(self):
+    self._use_block_diffusion()
+    data = _make_data()
+    data["targets_loss_mask"] = jnp.zeros_like(data["targets_loss_mask"])
+    model = _UniformNnxDecoder(self.config.vocab_size)
+    with self._cross_entropy_patch():
+      loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True)
+
+    self.assertEqual(int(aux["total_weights"]), 0)
+    self.assertEqual(float(aux["xent_sum"]), 0.0)
+    self.assertTrue(bool(jnp.isfinite(loss)))
+    self.assertEqual(float(loss), 0.0)
+
+  def test_causal_loss_rejects_block_bidirectional_attention(self):
+    self.config.attention_type = "block_diffusion"
+    model = _UniformNnxDecoder(self.config.vocab_size)
+
+    with self.assertRaisesRegex(ValueError, "would leak"):
+      pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True)
+
+  def test_block_diffusion_requires_all_explicit_masks(self):
+    self._use_block_diffusion()
+    model = _UniformNnxDecoder(self.config.vocab_size)
+
+    for mask_name in ("completion_mask", "corruption_mask", "targets_loss_mask"):
+      with self.subTest(mask_name=mask_name):
+        data = _make_data()
+        del data[mask_name]
+        with self.assertRaisesRegex(ValueError, mask_name):
+          pre_train.loss_fn(model, self.config, data, None, None, is_train=True)
+
+  def test_logit_alignment_matches_runtime_contract(self):
+    logits = jnp.arange(8, dtype=jnp.float32).reshape(1, 4, 2)
+
+    same_position = diffusion_scoring.align_logits_to_targets(logits, "same_position")
+    shifted = diffusion_scoring.align_logits_to_targets(logits, "shifted")
+
+    np.testing.assert_array_equal(same_position, logits)
+    np.testing.assert_array_equal(shifted, logits[:, [0, 0, 1, 2], :])
+
+  def test_shifted_alignment_uses_logical_positions_after_reordering(self):
+    positions = jnp.asarray([[0, 3, 1, 2]], dtype=jnp.int32)
+    logits = positions[..., None].astype(jnp.float32)
+
+    shifted = diffusion_scoring.align_logits_to_targets(
+        logits,
+        "shifted",
+        positions=positions,
+        validity_mask=jnp.ones_like(positions, dtype=jnp.bool_),
+    )
+
+    np.testing.assert_array_equal(shifted[..., 0], [[0, 2, 0, 1]])
+
+  def test_seeded_shifted_canvas_excludes_only_sequence_origin(self):
+    self.config.attention_type = "block_diffusion"
+    self.config.training_objective = "block_diffusion"
+    self.config.block_diffusion_logit_alignment = "shifted"
+    self.config.block_diffusion_canvas_policy = "seed_and_mask"
+    data = _make_data()
+    data["targets_loss_mask"] = jnp.ones_like(data["targets_segmentation"])
+    model = _UniformNnxDecoder(self.config.vocab_size)
+
+    with self._cross_entropy_patch():
+      _, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True)
+
+    self.assertEqual(int(aux["total_weights"]), 5)
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/unit/synthetic_data_test.py b/tests/unit/synthetic_data_test.py
index 9e29c04a15..f713037a13 100644
--- a/tests/unit/synthetic_data_test.py
+++ b/tests/unit/synthetic_data_test.py
@@ -15,6 +15,7 @@
 """Tests for synthetic data sharding."""
 
 import sys
+from types import SimpleNamespace
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -22,7 +23,7 @@
 from jax.experimental import mesh_utils
 
 from maxtext.configs import pyconfig
-from maxtext.input_pipeline.synthetic_data_processing import SyntheticDataIterator
+from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator, SyntheticDataIterator
 from maxtext.utils import sharding as maxtext_sharding
 from tests.utils.test_helpers import get_test_config_path
 
@@ -75,5 +76,25 @@ def test_synthetic_data_sharding(self, mesh_shape, shard_mode):
     self.assertEqual(inputs.sharding, expected_sharding)
 
 
+class PlaceHolderDataTest(absltest.TestCase):
+  """Tests placeholder batch metadata for non-loading hosts."""
+
+  def test_block_diffusion_placeholder_has_zero_loss_masks(self):
+    config = SimpleNamespace(
+        global_batch_size_to_load=jax.process_count(),
+        max_target_length=8,
+        training_objective="block_diffusion",
+    )
+
+    batch = next(PlaceHolderDataIterator.get_place_holder_synthetic_data(config))
+
+    self.assertIn("completion_mask", batch)
+    self.assertIn("corruption_mask", batch)
+    self.assertIn("targets_loss_mask", batch)
+    self.assertFalse(batch["completion_mask"].any())
+    self.assertFalse(batch["corruption_mask"].any())
+    self.assertFalse(batch["targets_loss_mask"].any())
+
+
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/unit/tokamax_splash_attention_mask_test.py b/tests/unit/tokamax_splash_attention_mask_test.py
index 590e8046b8..b1fe154029 100644
--- a/tests/unit/tokamax_splash_attention_mask_test.py
+++ b/tests/unit/tokamax_splash_attention_mask_test.py
@@ -20,6 +20,7 @@
 
 from maxtext.kernels.tokamax_splash_attention import splash_attention_mask
 from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info
+from maxtext.layers.attention_op import BlockCausalMask
 
 
 class TokamaxSplashAttentionMaskTest(absltest.TestCase):
@@ -56,6 +57,22 @@ def test_process_causal_mask_keeps_lazy_mask_function(self):
     self.assertIsNone(mask_info.partial_mask_blocks)
     np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32))
 
+  def test_process_block_causal_mask_keeps_lazy_mask_function(self):
+    mask = BlockCausalMask((8, 8), block_size=4)
+
+    mask_info, mask_function = splash_attention_mask_info.process_mask(
+        mask,
+        block_shape=(2, 2),
+        q_seq_shards=1,
+        kv_seq_shards=1,
+        return_dynamic_grid=True,
+    )
+
+    self.assertIs(mask_function, mask.mask_function)
+    self.assertIsNotNone(mask_info.q_sequence)
+    self.assertIsNone(mask_info.partial_mask_blocks)
+    np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32))
+
 
 if __name__ == "__main__":
   absltest.main()