diff --git a/src/maxtext/training_engine/README.md b/src/maxtext/training_engine/README.md
index 9793eec4f3..58b2797dff 100644
--- a/src/maxtext/training_engine/README.md
+++ b/src/maxtext/training_engine/README.md
@@ -296,6 +296,14 @@ and does not use the engine.
 | `model_scope(*inputs)` | none; the caller jits | inputs | yields the model and the placed inputs under the sharding rules |
 | `fwd_only(fn, *inputs)` | none; `fn` jits | inputs | `fn(model, ...)` |
 | `eval_step(batch)` | eval | params, rest, batch | loss, aux |
+| `prepare_fwd_bwd_input(micro_batch)` | none; `device_put` | micro-batch | a `PreparedFwdBwdInput`: the micro-batch mapped to loss inputs and placed on `fwd_bwd`'s input shardings, which `fwd_bwd` takes in its place |
+| `wait_for_device()` | none | | the seconds it blocked until the in-flight `fwd_bwd`s and `update` finished; dequeues nothing and writes no metrics |
+
+A caller may prepare the next step's first micro-batch between a step's last `fwd_bwd` and its
+`update()`, as Tunix does when it stages before a deferred update. Under `retry_when_tokens_dropped`
+the last micro-batch's overflow flag is then still pending: neither `prepare_fwd_bwd_input` nor
+`wait_for_device` resolves it, `update()` does, and a dropless rerun reads that micro-batch's own kept
+inputs, so the prepared micro-batch is unaffected (`maxtext_engine_stage_before_update_test.py`).
 
 Every micro-batch runs the same `fwd_bwd` program. The first micro-batch's gradients start the
 accumulator, and `accumulate` adds each later one's to it in place. The accumulator is not passed to
@@ -383,6 +391,8 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi
 | `maxtext_engine_model_build_test.py` | model construction with and without a mesh set by the caller |
 | `maxtext_engine_checkpoint_test.py` | resuming from a mid-step checkpoint through Orbax reproduces the uninterrupted step |
 | `maxtext_engine_grad_cast_test.py` | `cast_grads_after_all_reduce`: which gradients leave `fwd_bwd` uncast, and that moving their cast changes no number on CPU |
+| `maxtext_engine_prepared_input{,_dropless,_sharded}_test.py` | `prepare_fwd_bwd_input` and `wait_for_device`: a prepared micro-batch trains bit-identically on one device, through Tunix's `TrainerWorker` staging, under the dropless replay and final dropless programs, and on a four-device mesh with Tunix's payload; the wait leaves the in-flight queue and the metrics writes as they were |
+| `maxtext_engine_stage_before_update_test.py` | the order Tunix uses to stage before a deferred update (`fwd_bwd(last)`, `prepare_fwd_bwd_input(next)`, `wait_for_device()`, `update()`) under the dropless replay and final dropless programs, directly and through Tunix's `TrainerWorker`: bit-identical to the base order, including a rerun of the last micro-batch inside `update()` |
 | `maxtext_engine_{constructor,data_parallel,xaot,packing,profiling}_test.py`, `router_replay_engine_test.py` | construction, data parallelism, ahead-of-time compilation, sequence packing, profiling, router replay |
 
 ## Differences from pre_train/train.py
diff --git a/src/maxtext/training_engine/inflight_throttler.py b/src/maxtext/training_engine/inflight_throttler.py
index a809fe5e59..21fb91e012 100644
--- a/src/maxtext/training_engine/inflight_throttler.py
+++ b/src/maxtext/training_engine/inflight_throttler.py
@@ -65,6 +65,17 @@ def wait_for_next(self) -> None:
         self._flush_pending_metrics()
         self._pending_metrics = metrics
 
+  def block_until_ready(self) -> None:
+    """Waits for every inflight computation to finish, leaving the queue as it is.
+
+    Unlike `wait_for_all`, nothing is dequeued and no metrics are written: the next
+    `wait_for_next` / `wait_for_all` sees the same queue and finds it already finished.
+    """
+    with self._inflight_queue.mutex:
+      inflight = list(self._inflight_queue.queue)
+    for computation, _ in inflight:
+      jax.block_until_ready(computation)
+
   def wait_for_all(self) -> None:
     """Wait for all inflight computations to finish and log their metrics."""
     while not self._inflight_queue.empty():
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 049da90b6b..0f1d0ddc1b 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -470,6 +470,17 @@ def _to_aval(value: Any) -> Any:
   return jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=getattr(value, "sharding", None))
 
 
+@dataclasses.dataclass(frozen=True)
+class PreparedFwdBwdInput:
+  """A micro-batch already mapped to loss inputs and placed on the mesh.
+
+  Made by `MaxTextTrainingEngine.prepare_fwd_bwd_input`; `fwd_bwd` takes it in place of
+  the payload it came from.
+  """
+
+  batch: Any
+
+
 def _row_and_mask(payload: abstract_engine.RLTrainerPayload) -> tuple[jax.Array, jax.Array]:
   """Returns the `[batch, seq]` model row and its mask, prompt part then completion part.
 
@@ -2404,17 +2415,66 @@ def _start_running_sum(self, grads: Any) -> Any:
     cast = iter(self._cast_uncast_grads([g for g, is_uncast in zip(leaves, uncast) if is_uncast]))
     return treedef.unflatten([next(cast) if is_uncast else g for g, is_uncast in zip(leaves, uncast)])
 
+  def prepare_fwd_bwd_input(self, payload: abstract_engine.TrainerPayload) -> PreparedFwdBwdInput:
+    """Maps a micro-batch to loss inputs and places them on the mesh, ahead of its `fwd_bwd`.
+
+    For a caller that receives the micro-batch early (Tunix's `TrainerWorker.stage_fwd_bwd`,
+    while the device still computes the previous step's last micro-batch, and its update unless
+    the caller deferred it): the input mapping and the host-to-device copy run here, and
+    `fwd_bwd(prepared)` later dispatches without moving data. The traced inputs go to exactly
+    the shardings `fwd_bwd` compiles them onto (`_batch_data_shardings`), with the closed-over
+    half untouched, so the kernel, its inputs and the result are those of `fwd_bwd(payload)`.
+    Holds one micro-batch of inputs in device memory until that `fwd_bwd`.
+
+    Safe between a step's last `fwd_bwd` and its `update()`, the order Tunix uses to stage
+    before a deferred update. Under `retry_when_tokens_dropped` the last micro-batch's
+    overflow flag is still pending there: this neither reads nor resolves it, and the dropless
+    rerun `update()` may then run reads the inputs that micro-batch's pending check kept, never
+    the prepared ones. The mapping reads only the payload (no parameter, gradient or optimizer
+    state), so the prepared input is the same whether the update runs before or after it.
+
+    Args:
+      payload: Packed micro-batch training input.
+
+    Returns:
+      The input to pass to `fwd_bwd` instead of `payload`.
+    """
+    batch = self._prepare_batch(payload)
+    if self._mesh is not None and isinstance(batch, dict):
+      dynamic_batch, static_batch = _split_static_and_dynamic(batch)
+      dynamic_batch = jax.device_put(dynamic_batch, self._batch_data_shardings(dynamic_batch))
+      batch = {**dynamic_batch, **static_batch}
+    return PreparedFwdBwdInput(batch=batch)
+
+  def wait_for_device(self) -> float:
+    """Blocks until the computations this engine holds in flight have finished; returns the seconds waited.
+
+    For a caller that did work after dispatching and wants to know whether it hid behind the
+    device (Tunix's `TrainerWorker.stage_fwd_bwd` reports it): ~0 means the device was already
+    idle. It waits on what the in-flight throttler holds: each queued `fwd_bwd`'s loss and the
+    last `update`'s gradient norm, so the update counts. Anything dispatched before those was
+    waited on when it left the queue. A blocking wait, not an `is_ready()` poll: under Pathways
+    a readiness check is an RPC to the proxy, so a poll can read "not ready" for work that has
+    finished. Changes no state: the in-flight queue and its pending metrics are left as they
+    were (see `InflightThrottler.block_until_ready`), and so is a pending overflow check: under
+    `retry_when_tokens_dropped` the last micro-batch's flag, and its dropless rerun if it
+    overflowed, are still left to the next `fwd_bwd` or `update()`.
+    """
+    start = time.monotonic()
+    self._throttler.block_until_ready()
+    return time.monotonic() - start
+
   @_profiled_step("fwd_bwd", control_profile=True)
-  def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> None:
+  def fwd_bwd(self, payload: abstract_engine.TrainerPayload | PreparedFwdBwdInput, **kwargs: Any) -> None:
     """Executes a micro-batch forward-backward pass and accumulates gradients.
 
     Args:
-      payload: Packed micro-batch training input.
+      payload: Packed micro-batch training input, or what `prepare_fwd_bwd_input` made of it.
       **kwargs: Implementation-specific options. `apply_optimizer` (bool), which Tunix passes as the
         orchestrator's flag for the step's final micro-batch, selects the final micro-batch's dropless program
         under `retry_dropless_final_micro_batch`; everything else is ignored.
     """
-    batch = self._prepare_batch(payload)
+    batch = payload.batch if isinstance(payload, PreparedFwdBwdInput) else self._prepare_batch(payload)
     apply_optimizer = kwargs.get("apply_optimizer")
     final = False
     if self._dropless_replay and self._final_dropless:
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
new file mode 100644
index 0000000000..9db43e0f33
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
@@ -0,0 +1,172 @@
+# 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
+#
+#    https://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.
+"""prepare_fwd_bwd_input under the dropless replay and the final micro-batch's dropless program.
+
+The team's MLPerf trainer configuration trains with `retry_when_tokens_dropped` (a micro-batch whose MoE ragged buffer
+overflowed is accumulated as zero and rerun dropless, from the inputs it ran on); `retry_dropless_final_micro_batch` (a
+step's last micro-batch runs the dropless program directly) is an option on top of it. A micro-batch prepared ahead must
+take both paths exactly like the payload it came from: the overflowing micro-batch is replayed from its prepared inputs,
+the final one runs the final program, and the weights, losses and replay counts are bit-identical.
+
+The rig is `maxtext_engine_final_dropless_test.py`'s: a one-layer Qwen3.5 MoE on CPU whose loss reports the overflow
+flag RoutedMoE would, from which program its graph is, and scales the loss per program so the kernels can be told
+apart.
+"""
+# pylint: disable=protected-access
+
+import os
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from maxtext.models import models
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import pytest
+
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg, _payload, _MARKED
+from tests.post_training.unit.maxtext_engine_final_dropless_test import _program
+
+pytestmark = [pytest.mark.post_training]
+
+# Per-program loss multipliers, so the losses say which kernel ran each micro-batch.
+_SCALES = {"regular": 1.0, "replay": 2.0, "final": 3.0}
+
+# Two steps of three micro-batches. Step 1's first micro-batch overflows, so the one prepared ahead is replayed.
+_STEPS = ((10, _MARKED, 30), (_MARKED, 20, 30))
+
+# Where the next step's first micro-batch is prepared, relative to the step's update.
+_BEFORE_UPDATE = "before_update"
+_AFTER_UPDATE = "after_update"  # where Tunix's stage lands when the update runs inside the last train call
+
+
+class PreparedInputDroplessTest(unittest.TestCase):
+
+  def setUp(self):
+    os.environ["NEW_MODEL_DESIGN"] = "1"
+    os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+
+  def _engine(self):
+    """An engine on the one-layer MoE rig, with the dropless replay and the final dropless program on."""
+    cfg = _cfg(override_logical_axis_rules=True, ragged_buffer_factor=2.5)
+    mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    real_model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
+    with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=real_model):
+      engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh)
+    # The ring-of-experts prerequisites of retry_when_tokens_dropped cannot run on CPU, so the switches are set
+    # directly, as the dropless replay and final-dropless tests do.
+    engine._dropless_replay = True
+    engine._final_dropless = True
+    base_loss = maxtext_engine.make_router_replay_loss_fn(cfg)
+
+    def loss_fn(model, **batch):
+      program = _program(model)
+      loss, aux = base_loss(model, **batch)
+      aux = {k: v for k, v in aux.items() if k != "has_moe_overflow"}
+      aux["has_moe_overflow"] = jnp.logical_and(batch["inputs"][0, 0] == _MARKED, program == "regular")
+      aux["xent_sum"] = aux["xent_sum"] * _SCALES[program]
+      return loss, aux
+
+    engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn)
+    engine.with_loss_fn(loss_fn, has_aux=True)
+    engine.compile(None)  # As Tunix's TrainerWorker does at bring-up: the kernels compile on the first fwd_bwd.
+    return engine
+
+  def _train(self, engine, prepare_ahead: bool, order: str = _AFTER_UPDATE, prepare_first: bool = False):
+    """Runs `_STEPS`; returns the final parameters, per-step metrics and the prepared micro-batches' replays.
+
+    With `prepare_ahead`, each step's first micro-batch after the first step is prepared right after the previous
+    step's last `fwd_bwd` (`apply_optimizer=True`, so it ran the final program), before or after the update, and
+    `wait_for_device` follows, as Tunix's `TrainerWorker.stage_fwd_bwd` does.
+    """
+    prepared = None
+    if prepare_first:
+      prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[0][0]))
+    metrics = []
+    for step, tokens in enumerate(_STEPS):
+      for j, first_token in enumerate(tokens):
+        final = j == len(tokens) - 1
+        if j == 0 and prepared is not None:
+          self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+          engine.fwd_bwd(prepared, apply_optimizer=False)
+          prepared = None
+        else:
+          engine.fwd_bwd(_payload(first_token), apply_optimizer=final)
+      has_next = step + 1 < len(_STEPS)
+      if prepare_ahead and has_next and order == _BEFORE_UPDATE:
+        prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+      with mock.patch.object(engine, "record_metrics", wraps=engine.record_metrics) as record:
+        engine.update()
+      recorded = {call.args[0]: call.args[1] for call in record.call_args_list}
+      if prepare_ahead and has_next and order == _AFTER_UPDATE:
+        prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append(
+          {
+              "replays": int(recorded["moe_dropless_replays"]),
+              "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+              "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+          }
+      )
+    self.assertIsNone(prepared)
+    params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+    return params, metrics
+
+  def _assert_same_training(self, ref, new):
+    ref_params, ref_metrics = ref
+    new_params, new_metrics = new
+    self.assertEqual(ref_metrics, new_metrics)
+    self.assertEqual(len(ref_params), len(new_params))
+    for r, n in zip(ref_params, new_params):
+      np.testing.assert_array_equal(r, n)
+
+  def test_a_prepared_micro_batch_is_replayed_and_trains_bit_identically(self):
+    reference = self._train(self._engine(), prepare_ahead=False)
+    # Not vacuous: one replay per step (the marked micro-batches), and step 1's is the prepared one.
+    self.assertEqual([m["replays"] for m in reference[1]], [1, 1])
+    for order in (_AFTER_UPDATE, _BEFORE_UPDATE):
+      with self.subTest(order=order):
+        engine = self._engine()
+        with mock.patch.object(engine, "_compile_for_batch", wraps=engine._compile_for_batch) as compiles:
+          prepared = self._train(engine, prepare_ahead=True, order=order)
+        self._assert_same_training(reference, prepared)
+        self.assertEqual(compiles.call_count, 1)  # no kernel was rebuilt for the prepared input
+
+  def test_the_first_micro_batch_of_the_run_can_be_prepared(self):
+    # Tunix never stages a run's first micro-batch, but the background compiles that start from it must still
+    # produce the kernels a payload would.
+    reference = self._train(self._engine(), prepare_ahead=False)
+    prepared = self._train(self._engine(), prepare_ahead=True, prepare_first=True)
+    self._assert_same_training(reference, prepared)
+
+  def test_the_comparison_sees_a_different_prepared_micro_batch(self):
+    # Control: preparing the wrong micro-batch (step 1's second instead of its first) is caught.
+    reference = self._train(self._engine(), prepare_ahead=False)
+    engine = self._engine()
+    real_prepare = engine.prepare_fwd_bwd_input
+    with mock.patch.object(
+        engine, "prepare_fwd_bwd_input", side_effect=lambda payload: real_prepare(_payload(_STEPS[1][1]))
+    ):
+      wrong = self._train(engine, prepare_ahead=True)
+    self.assertNotEqual(reference[1], wrong[1])
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
new file mode 100644
index 0000000000..3ac9365492
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
@@ -0,0 +1,271 @@
+# 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
+#
+#    https://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.
+
+"""prepare_fwd_bwd_input with Tunix's real model input, on a four-device mesh.
+
+`maxtext_engine_prepared_input_test` uses a one-device mesh and a flat toy input. Production feeds Tunix's
+`_algo_model_input`: {"train_example": RLTrainerPayload (a flax.struct pytree with static `metadata`),
+"algo_config": a closed-over object, "pad_id", "eos_id"}, on a mesh whose batch axis shards. Here the batch is split
+over four host devices, and a prepared micro-batch, with `wait_for_device` after it as Tunix's
+`TrainerWorker.stage_fwd_bwd` does, must train bit-identically to the payload it came from.
+
+Everything runs in a subprocess, for the reason `maxtext_engine_data_parallel_test.py` gives: the CPU backend reads
+`--xla_force_host_platform_device_count` only when it initializes, which has already happened by the time pytest
+imports this file, so setting the flag here would leave one device and skip the tests green.
+`test_prepared_input_on_a_four_device_cpu_mesh` re-execs the module with the flag appended and refuses to pass unless
+the child reports tests actually run.
+"""
+# pylint: disable=protected-access
+
+import dataclasses
+import functools
+import os
+import re
+import subprocess
+import sys
+import types
+import unittest
+from unittest import mock
+
+from absl import flags
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import optax
+import pytest
+
+from tests.utils.test_helpers import get_test_config_path
+
+# training_engine imports tunix, so these tests need the post-training dependency bundle.
+pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only]
+
+_REQUIRED_DEVICES = 4
+_SENTINEL = "MAXTEXT_ENGINE_PREPARED_INPUT_SHARDED_TESTS_PASSED"
+_RAN = re.compile(rf"{_SENTINEL} ran=(\d+)")
+
+B, P, C, L, K = 4, 8, 16, 2, 3
+
+
+@pytest.mark.post_training
+@pytest.mark.cpu_only
+def test_prepared_input_on_a_four_device_cpu_mesh():
+  """The only test pytest collects here; the class below runs inside the child process."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = f"{env.get('XLA_FLAGS', '')} --xla_force_host_platform_device_count={_REQUIRED_DEVICES}".strip()
+  env["JAX_PLATFORMS"] = "cpu"
+  # The child imports `tests.utils.test_helpers`, which pytest puts on the path for us and a bare interpreter does
+  # not.
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+
+  report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert result.returncode == 0, report
+  ran = _RAN.search(result.stdout)
+  # An exit status of 0 is also what a run that skipped everything produces.
+  assert ran, f"the child did not report a completed run\n{report}"
+  assert int(ran.group(1)) > 0, f"every test in the child skipped\n{report}"
+
+
+class _Model(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0, 0.5]))
+
+
+def _payload(seed: int):
+  from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
+
+  rng = np.random.default_rng(seed)
+  return datatypes.RLTrainerPayload(
+      prompt_ids=rng.integers(1, 50, (B, P)).astype(np.int32),
+      prompt_mask=np.ones((B, P), np.float32),
+      completion_ids=rng.integers(1, 50, (B, C)).astype(np.int32),
+      completion_mask=(rng.random((B, C)) > 0.2).astype(np.float32),
+      advantages=rng.normal(size=(B,)).astype(np.float32),
+      old_per_token_logps=-rng.random((B, C)).astype(np.float32),
+      routed_experts=rng.integers(0, 8, (B, P + C, L, K)).astype(np.int16),
+      metadata={"trajectory_ids": tuple(f"t{seed}_{i}" for i in range(B))},
+  )
+
+
+def _loss_fn(model, train_example, algo_config, pad_id, eos_id):
+  del eos_id
+  w = model.weights[...]
+  x = train_example.completion_ids.astype(jnp.float32) / 50.0
+  m = train_example.completion_mask * (train_example.completion_ids != pad_id)
+  routed = jnp.mean(train_example.routed_experts.astype(jnp.float32), (2, 3))
+  pred = w[0] * x + w[1] * x**2 * algo_config.scale + w[2] * routed[:, P:] / 8.0 + train_example.old_per_token_logps
+  per_tok = pred * train_example.advantages[:, None] * m
+  return (
+      abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(per_tok), denominator=jnp.sum(m)),
+      {},
+  )
+
+
+class NestedPayloadPreparedInputTest(absltest.TestCase):
+
+  __test__ = False  # collected only via the subprocess entry point at the top of this file.
+
+  def setUp(self):
+    """A four-device mesh and a toy three-weight model, built without a checkpoint."""
+    super().setUp()
+    if jax.device_count() < _REQUIRED_DEVICES:
+      self.skipTest(f"needs {_REQUIRED_DEVICES} devices; run through pytest, which re-execs with them")
+    self.config = pyconfig.initialize(
+        [None, get_test_config_path()],
+        model_name="llama3.1-8b",
+        run_name="test_run",
+        base_output_directory=self.create_tempdir().full_path,
+        init_weights_seed=42,
+        micro_batch_size_to_train_on=B,
+        gradient_accumulation_steps=1,
+        enable_dropout=False,
+        record_internal_nn_metrics=False,
+        enable_tensorboard=False,
+        tensorboard_dir=self.create_tempdir().full_path,
+        skip_jax_distributed_system=True,
+        enable_checkpointing=False,
+        profiler_steps=0,
+    )
+    self.mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(self.config), self.config.mesh_axes)
+    for patcher in (
+        mock.patch.object(
+            maxtext_engine.train_utils,
+            "create_training_optimizer",
+            return_value=(lambda step: jnp.array(0.001), optax.adam(0.05)),
+        ),
+        mock.patch.object(
+            maxtext_engine.model_creation_utils,
+            "from_pretrained",
+            side_effect=lambda *a, **k: (_Model(), self.mesh),
+        ),
+    ):
+      self.addCleanup(patcher.stop)
+      patcher.start()
+    self.algo_config = types.SimpleNamespace(scale=0.75)
+
+  def _engine(self):
+    """An engine fed through Tunix's `_algo_model_input`, compiled on the first payload."""
+    from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=import-outside-toplevel
+
+    engine = maxtext_engine.MaxTextTrainingEngine(self.config)
+    engine.with_loss_fn(_loss_fn, has_aux=True)
+    engine.with_gen_model_input_fn(
+        functools.partial(
+            algorithm_adapter._algo_model_input,
+            algo_config=self.algo_config,
+            pad_id=0,
+            eos_id=1,
+        )
+    )
+    engine.compile(_payload(0))
+    return engine
+
+  def _train(self, engine, payloads, prepare_at):
+    """Two steps of two micro-batches; those in `prepare_at` prepared ahead, after the previous update."""
+    prepared = {}
+    if 0 in prepare_at:
+      prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
+      engine.wait_for_device()
+    metrics = []
+    for step in range(2):
+      for j in range(2):
+        i = 2 * step + j
+        engine.fwd_bwd(prepared.pop(i) if i in prepared else payloads[i])
+      engine.update()
+      if 2 * step + 2 in prepare_at:
+        # Prepared while the update is still on the device, as stage_fwd_bwd does, which then waits for the device
+        # to report whether that hid.
+        prepared[2 * step + 2] = engine.prepare_fwd_bwd_input(payloads[2 * step + 2])
+        engine.wait_for_device()
+      buf = engine.get_metrics(clear_cache=True)
+      metrics.append(np.asarray(buf.weighted_metrics["loss"].compute()).tolist())
+    return np.asarray(engine.model.weights[...]), metrics
+
+  def test_bit_identical_on_a_sharded_mesh(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    ref, new = self._engine(), self._engine()
+    with (
+        mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+        mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+    ):
+      ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
+      new_w, new_m = self._train(new, payloads, prepare_at={0, 2})
+    np.testing.assert_array_equal(ref_w, new_w)
+    self.assertEqual(ref_m, new_m)
+    self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+
+  def test_prepared_leaves_are_sharded_like_the_compile(self):
+    engine = self._engine()
+    prepared = engine.prepare_fwd_bwd_input(_payload(2))
+    dynamic, static = maxtext_engine._split_static_and_dynamic(prepared.batch)
+    self.assertEqual(set(dynamic), {"train_example"})
+    self.assertEqual(set(static), {"algo_config", "pad_id", "eos_id"})
+    shardings = engine._batch_data_shardings(dynamic)
+    example, sharding = dynamic["train_example"], shardings["train_example"]
+    self.assertEqual(example.metadata, {})  # stripped by _prepare_batch
+    split_2d = 0
+    for field in (
+        "completion_ids",
+        "completion_mask",
+        "old_per_token_logps",
+        "routed_experts",
+        "advantages",
+        "prompt_ids",
+    ):
+      leaf, want = getattr(example, field), getattr(sharding, field)
+      self.assertIsInstance(leaf, jax.Array, field)
+      self.assertEqual(leaf.sharding, want, field)
+      rows = {s.index[0] for s in leaf.addressable_shards}
+      if leaf.ndim >= 2 and len(rows) == _REQUIRED_DEVICES:
+        split_2d += 1
+    # Not vacuous: the 2-D leaves really are split four ways.
+    self.assertGreaterEqual(split_2d, 3)
+
+  def test_one_token_change_is_seen(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    changed = list(payloads)
+    ids = np.array(payloads[2].completion_ids)
+    ids[0, 0] = ids[0, 0] % 49 + 1
+    changed[2] = dataclasses.replace(payloads[2], completion_ids=ids)
+    ref_w, _ = self._train(self._engine(), payloads, prepare_at=set())
+    new_w, _ = self._train(self._engine(), changed, prepare_at={2})
+    self.assertFalse(np.array_equal(ref_w, new_w))
+
+
+_SUITE = (NestedPayloadPreparedInputTest,)
+
+
+if __name__ == "__main__":
+  if jax.device_count() < _REQUIRED_DEVICES:
+    raise SystemExit(
+        f"needs {_REQUIRED_DEVICES} devices, got {jax.device_count()}; run this through pytest, which sets "
+        f"XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}"
+    )
+  flags.FLAGS(sys.argv[:1])  # absltest's create_tempdir reads --test_tmpdir; a bare unittest run parses no flags.
+  _loader = unittest.defaultTestLoader
+  _result = unittest.TextTestRunner(verbosity=2).run(
+      unittest.TestSuite(_loader.loadTestsFromTestCase(cls) for cls in _SUITE)
+  )
+  if not _result.wasSuccessful():
+    sys.exit(1)
+  print(f"{_SENTINEL} ran={_result.testsRun - len(_result.skipped)}")
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
new file mode 100644
index 0000000000..6552caf720
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
@@ -0,0 +1,387 @@
+# 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
+#
+#    https://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 MaxTextTrainingEngine.prepare_fwd_bwd_input and wait_for_device.
+
+A micro-batch prepared ahead (mapped and placed on the mesh while earlier work
+still runs) must train exactly like the payload it came from: same compiled
+kernel, same inputs, bit-identical weights and metrics. Waiting for the device
+afterwards (what Tunix's `TrainerWorker.stage_fwd_bwd` does, to report whether
+the transfer hid behind the device) must change nothing either.
+
+`maxtext_engine_prepared_input_dropless_test.py` covers the dropless replay and
+the final micro-batch's dropless program, and
+`maxtext_engine_prepared_input_sharded_test.py` a four-device mesh with Tunix's
+real model input.
+"""
+# pylint: disable=protected-access
+
+import dataclasses
+import types
+from typing import Any
+from unittest import mock
+
+from absl.testing import absltest
+from flax import nnx
+from flax import struct
+import jax
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import inflight_throttler
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+import numpy as np
+import optax
+import pytest
+
+# training_engine imports tunix, so these tests need the post-training dependency bundle.
+pytestmark = [pytest.mark.post_training]
+
+
+class _Model(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0]))
+
+
+@struct.dataclass(frozen=True, kw_only=True)
+class _Payload(abstract_engine.TrainerPayload):
+  token_ids: Any
+  token_mask: Any
+  metadata: dict[str, Any] = struct.field(pytree_node=False, default_factory=dict)
+
+
+def _payload(seed: int) -> _Payload:
+  rng = np.random.default_rng(seed)
+  return _Payload(
+      token_ids=rng.normal(size=(2, 4)).astype(np.float32),
+      token_mask=(rng.random((2, 4)) > 0.25).astype(np.float32),
+      metadata={"seed": seed},
+  )
+
+
+def _loss_fn(model, tokens, mask, scale):
+  # Depends on the data and on both weights, so any change to either shows.
+  pred = model.weights[...][0] * tokens + model.weights[...][1] * tokens**2
+  return (
+      abstract_engine.WeightedMetric(
+          unreduced_sum=jnp.sum(pred * mask) * scale,
+          denominator=jnp.sum(mask),
+      ),
+      {},
+  )
+
+
+def _model_input(payload: _Payload) -> dict[str, Any]:
+  # Tunix's shape: traced arrays plus a closed-over Python constant.
+  return {"tokens": payload.token_ids, "mask": payload.token_mask, "scale": 2.0}
+
+
+# Where a step's next first micro-batch is prepared, relative to the step's update.
+_BEFORE_UPDATE = "before_update"  # a separate update call follows the last fwd_bwd
+_AFTER_UPDATE = "after_update"  # the update runs inside the last fwd_bwd call (google/tunix#2719)
+
+
+class PrepareFwdBwdInputTest(absltest.TestCase):
+
+  def setUp(self):
+    """A one-device mesh and a toy two-weight model, built without a checkpoint."""
+    super().setUp()
+    self.config = pyconfig.initialize(
+        [None, get_test_config_path()],
+        model_name="llama3.1-8b",
+        run_name="test_run",
+        base_output_directory=self.create_tempdir().full_path,
+        init_weights_seed=42,
+        micro_batch_size_to_train_on=2,
+        gradient_accumulation_steps=1,
+        enable_dropout=False,
+        record_internal_nn_metrics=False,
+        enable_tensorboard=False,
+        tensorboard_dir=self.create_tempdir().full_path,
+        skip_jax_distributed_system=True,
+        enable_checkpointing=False,
+        profiler_steps=0,
+    )
+    self.mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(self.config), self.config.mesh_axes)
+    for patcher in (
+        mock.patch.object(
+            maxtext_engine.train_utils,
+            "create_training_optimizer",
+            return_value=(lambda step: jnp.array(0.001), optax.adam(0.05)),
+        ),
+        mock.patch.object(
+            maxtext_engine.model_creation_utils,
+            "from_pretrained",
+            side_effect=lambda *a, **k: (_Model(), self.mesh),
+        ),
+    ):
+      self.addCleanup(patcher.stop)
+      patcher.start()
+
+  def _engine(self, compiled: bool) -> maxtext_engine.MaxTextTrainingEngine:
+    engine = maxtext_engine.MaxTextTrainingEngine(self.config)
+    engine.with_loss_fn(_loss_fn)
+    engine.with_gen_model_input_fn(_model_input)
+    if compiled:
+      engine.compile(_payload(0))
+    return engine
+
+  def _train(self, engine, payloads, prepare_at: set[int], wait: bool = False, order: str = _BEFORE_UPDATE):
+    """Two steps of two micro-batches; micro-batches in `prepare_at` go in prepared.
+
+    A step's next first micro-batch is prepared right after the step's last `fwd_bwd`, while that micro-batch
+    and the update are still on the device: before the update with `order=_BEFORE_UPDATE`, after it with
+    `_AFTER_UPDATE` (where Tunix's `TrainerWorker.stage_fwd_bwd` lands once the update runs inside the last
+    train call). With `wait`, `wait_for_device` follows each prepare, as it does there too.
+    """
+    prepared = {}
+
+    def prepare(i):
+      if i not in prepare_at or i >= len(payloads):
+        return
+      prepared[i] = engine.prepare_fwd_bwd_input(payloads[i])
+      if wait:
+        queued = engine._throttler._inflight_queue.qsize()
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+        # The in-flight queue is left for the next dispatch, or a drain, to empty as before.
+        self.assertEqual(engine._throttler._inflight_queue.qsize(), queued)
+
+    prepare(0)
+    metrics = []
+    for step in range(2):
+      for j in range(2):
+        i = 2 * step + j
+        engine.fwd_bwd(prepared.pop(i) if i in prepared else payloads[i])
+      if order == _BEFORE_UPDATE:
+        prepare(2 * step + 2)
+      engine.update()
+      if order == _AFTER_UPDATE:
+        prepare(2 * step + 2)
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append({k: np.asarray(v).tolist() for k, v in sorted(buffer.scalar_metrics.items())})
+      metrics[-1]["loss"] = np.asarray(buffer.weighted_metrics["loss"].compute()).tolist()
+    self.assertEqual(prepared, {})
+    return np.asarray(engine.model.weights[...]), metrics
+
+  def test_prepared_micro_batches_train_bit_identically(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    for compiled in (True, False):
+      for order in (_BEFORE_UPDATE, _AFTER_UPDATE):
+        with self.subTest(compiled=compiled, order=order):
+          ref = self._engine(compiled)
+          new = self._engine(compiled)
+          with (
+              mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+              mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+          ):
+            ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
+            new_w, new_m = self._train(new, payloads, prepare_at={0, 2}, order=order)
+          np.testing.assert_array_equal(ref_w, new_w)
+          self.assertEqual(ref_m, new_m)
+          # No kernel was rebuilt for a prepared input.
+          self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+          self.assertEqual(new.train_step, 2)
+
+  def test_waiting_for_the_device_changes_nothing(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    for compiled in (True, False):
+      for order in (_BEFORE_UPDATE, _AFTER_UPDATE):
+        with self.subTest(compiled=compiled, order=order):
+          ref = self._engine(compiled)
+          new = self._engine(compiled)
+          with (
+              mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+              mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+              mock.patch.object(new._metrics_logger, "write_metrics", wraps=new._metrics_logger.write_metrics) as new_w,
+              mock.patch.object(ref._metrics_logger, "write_metrics", wraps=ref._metrics_logger.write_metrics) as ref_w,
+          ):
+            ref_weights, ref_m = self._train(ref, payloads, prepare_at=set())
+            new_weights, new_m = self._train(new, payloads, prepare_at={0, 2}, wait=True, order=order)
+          np.testing.assert_array_equal(ref_weights, new_weights)
+          self.assertEqual(ref_m, new_m)
+          self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+          # The same metrics writes, in the same order (the wait writes none).
+          self.assertEqual(
+              [c.args[0].id for c in ref_w.call_args_list],
+              [c.args[0].id for c in new_w.call_args_list],
+          )
+
+  def test_wait_for_device_waits_on_every_queued_computation_and_keeps_the_queue(self):
+    engine = self._engine(True)
+    engine.fwd_bwd(_payload(1))
+    engine.fwd_bwd(_payload(2))
+    engine.update()
+    queue = engine._throttler._inflight_queue
+    queued = list(queue.queue)
+    # The update queues its gradient norm with the step's metrics, after the micro-batches' losses.
+    self.assertIsNotNone(queued[-1][1])
+    waited_on = []
+    block = jax.block_until_ready
+
+    def spy(x):
+      waited_on.append(x)
+      return block(x)
+
+    with mock.patch.object(inflight_throttler.jax, "block_until_ready", side_effect=spy):
+      waited = engine.wait_for_device()
+    self.assertIsInstance(waited, float)
+    self.assertGreaterEqual(waited, 0.0)
+    # Nothing dequeued: the next dispatch or drain finds the same queue.
+    self.assertEqual(list(queue.queue), queued)
+    # Every queued computation was waited on, in order: each micro-batch's loss and the update's gradient norm,
+    # so a caller's wait covers the update.
+    self.assertEqual([id(x) for x in waited_on], [id(leaves) for leaves, _ in queued])
+    token = getattr(engine, "last_update_token", None)
+    if token is not None:
+      # With the readiness tokens of AI-Hypercomputer/maxtext#5571 the update's token is that same norm, so a
+      # caller that waits on it after `wait_for_device` (Tunix's `TrainerWorker.stage_fwd_bwd`) does not wait again.
+      self.assertIs(queued[-1][0][0], token)
+
+  def test_the_comparison_sees_a_one_token_difference(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    changed = list(payloads)
+    tokens = np.array(payloads[2].token_ids)
+    tokens[0, 0] += 1e-3
+    changed[2] = dataclasses.replace(payloads[2], token_ids=tokens)
+    ref_w, _ = self._train(self._engine(True), payloads, prepare_at=set())
+    new_w, _ = self._train(self._engine(True), changed, prepare_at={2})
+    self.assertFalse(np.array_equal(ref_w, new_w))
+
+  def test_staged_through_tunix_trainer_worker_trains_bit_identically(self):
+    """The path Tunix's pipelined step boundary takes, through a real `TrainerWorker`.
+
+    Each step's last micro-batch runs the update in the same call (`fwd_bwd(apply_optimizer=True)`,
+    google/tunix#2719). With staging, that call returns at the update's dispatch (`wait_for_update=False`) and the
+    next step's first micro-batch is staged right after it: `stage_fwd_bwd` calls `prepare_fwd_bwd_input` and
+    `wait_for_device`, and `fwd_bwd_staged` trains what it returned.
+    """
+    from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
+    from tunix.experimental.worker import trainer_worker  # pylint: disable=import-outside-toplevel
+
+    if not hasattr(trainer_worker.TrainerWorker, "stage_fwd_bwd"):
+      self.skipTest("this Tunix has no TrainerWorker.stage_fwd_bwd (its pipelined step boundary)")
+
+    payloads = [_payload(s) for s in (1, 2, 3, 4, 5, 6)]
+
+    def run(stage: bool):
+      engine = self._engine(True)
+      worker = trainer_worker.TrainerWorker(trainer_factory=lambda: engine)
+      worker.initialize()
+      with (
+          mock.patch.object(engine, "prepare_fwd_bwd_input", wraps=engine.prepare_fwd_bwd_input) as prepare,
+          mock.patch.object(engine, "wait_for_device", wraps=engine.wait_for_device) as wait,
+      ):
+        staged_id = None
+        for step in range(3):
+          for j in range(2):
+            i = 2 * step + j
+            kwargs = {"apply_optimizer": True, "wait_for_update": not stage} if j == 1 else {}
+            if staged_id is not None:
+              resp = worker.fwd_bwd_staged(staged_id, **kwargs)
+              staged_id = None
+            else:
+              resp = worker.fwd_bwd(datatypes.TrainRequest(request_id=f"r{i}", payload=payloads[i]), **kwargs)
+            if j == 1:
+              self.assertEqual(resp.metadata["train_step"], step + 1)
+              if stage and i + 1 < len(payloads):
+                staged = worker.stage_fwd_bwd(datatypes.TrainRequest(request_id=f"r{i + 1}", payload=payloads[i + 1]))
+                # The engine's wait_for_device answered how long the device still ran.
+                self.assertGreaterEqual(staged.metadata["device_wait_s"], 0.0)
+                staged_id = f"r{i + 1}"
+      return np.asarray(engine.model.weights[...]), prepare.call_count, wait.call_count
+
+    ref_w, ref_prepared, ref_waits = run(stage=False)
+    new_w, new_prepared, new_waits = run(stage=True)
+    np.testing.assert_array_equal(ref_w, new_w)
+    self.assertEqual((ref_prepared, new_prepared), (0, 2))
+    self.assertEqual((ref_waits, new_waits), (0, 2))
+
+  def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
+    engine = self._engine(True)
+    engine.fwd_bwd(_payload(1))
+    before = (
+        engine.micro_step_count,
+        engine.train_step,
+        np.asarray(engine.model.weights[...]),
+    )
+    prepared = engine.prepare_fwd_bwd_input(_payload(2))
+    self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+    self.assertEqual(prepared.batch["scale"], 2.0)
+    dynamic, _ = maxtext_engine._split_static_and_dynamic(prepared.batch)
+    shardings = engine._batch_data_shardings(dynamic)
+    for key in ("tokens", "mask"):
+      self.assertIsInstance(prepared.batch[key], jax.Array)
+      self.assertEqual(prepared.batch[key].sharding, shardings[key])
+    np.testing.assert_array_equal(np.asarray(prepared.batch["tokens"]), _payload(2).token_ids)
+    after = (
+        engine.micro_step_count,
+        engine.train_step,
+        np.asarray(engine.model.weights[...]),
+    )
+    self.assertEqual(before[:2], after[:2])
+    np.testing.assert_array_equal(before[2], after[2])
+
+
+class _Metrics:
+
+  def __init__(self):
+    self.written = []
+
+  def write_metrics(self, metrics, **kwargs):
+    del kwargs
+    self.written.append(metrics)
+
+
+class InflightThrottlerBlockUntilReadyTest(absltest.TestCase):
+
+  def _throttler(self):
+    logger = _Metrics()
+    config = types.SimpleNamespace(max_inflight_computations=2)
+    return inflight_throttler.InflightThrottler(config, logger), logger
+
+  def test_block_until_ready_neither_dequeues_nor_writes(self):
+    throttler, logger = self._throttler()
+    throttler.add_computation(jnp.ones(3) * 2, "m1")
+    throttler.add_computation(jnp.ones(3) * 3, None)
+    throttler.block_until_ready()
+    self.assertEqual(throttler._inflight_queue.qsize(), 2)
+    self.assertEqual(logger.written, [])
+    for leaves, _ in list(throttler._inflight_queue.queue):
+      self.assertTrue(all(leaf.is_ready() for leaf in leaves))
+
+  def test_a_drain_after_it_writes_what_it_would_have(self):
+    written = []
+    for block_first in (False, True):
+      throttler, logger = self._throttler()
+      throttler.add_computation(jnp.ones(3), "m1")
+      throttler.add_computation(jnp.ones(3), "m2")
+      if block_first:
+        throttler.block_until_ready()
+      throttler.wait_for_all()
+      written.append(logger.written)
+      self.assertTrue(throttler._inflight_queue.empty())
+    self.assertEqual(written[0], ["m1", "m2"])
+    self.assertEqual(written[0], written[1])
+
+  def test_block_until_ready_on_an_empty_queue_returns(self):
+    throttler, logger = self._throttler()
+    throttler.block_until_ready()
+    self.assertEqual(logger.written, [])
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
new file mode 100644
index 0000000000..fa9090c82a
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
@@ -0,0 +1,260 @@
+# 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
+#
+#    https://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.
+"""The order Tunix uses when it stages a step's first micro-batch before the previous step's update.
+
+Tunix's staged step boundary (`stage_before_update`) calls, at a staged boundary: `fwd_bwd(last, apply_optimizer=True)`,
+then `prepare_fwd_bwd_input(next step's first micro-batch)`, then `wait_for_device()`, then `update()`, and later
+`fwd_bwd(prepared)`. Under `retry_when_tokens_dropped` without `retry_dropless_final_micro_batch` (the team's MLPerf
+trainer configuration sets the first; the second is opt-in), the step's last micro-batch can overflow: its flag is still
+pending when the next micro-batch is prepared, and its dropless rerun runs inside `update()`, after the prepare. The
+prepared micro-batch must come out of that untouched, and the whole run must train bit-identically to the base order,
+`fwd_bwd(last, apply_optimizer=True)`, `update()`, `fwd_bwd(next payload)`.
+
+The rig is `maxtext_engine_prepared_input_dropless_test.py`'s (a one-layer Qwen3.5 MoE on CPU whose loss reports the
+overflow flag RoutedMoE would), with the final micro-batch's dropless program off.
+"""
+# pylint: disable=protected-access
+
+import dataclasses
+import os
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from maxtext.models import models
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import pytest
+
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg, _payload, _MARKED
+from tests.post_training.unit.maxtext_engine_final_dropless_test import _program
+
+pytestmark = [pytest.mark.post_training]
+
+# Per-program loss multipliers, so the losses say which kernel ran each micro-batch.
+_SCALES = {"regular": 1.0, "replay": 2.0, "final": 3.0}
+
+# A second overflowing content, so a rerun can be told from the micro-batch prepared while it is pending.
+_MARKED_B = 60
+
+# Three steps of three micro-batches. Every step's last micro-batch overflows. At both staged boundaries (after steps
+# 0 and 1) its rerun runs inside `update()` with the next step's first micro-batch already prepared; after step 2 it
+# runs inside the last `update()`. Step 1's first (prepared) micro-batch overflows too.
+_STEPS = ((10, 20, _MARKED), (_MARKED_B, 20, _MARKED), (40, 20, _MARKED))
+
+
+class StageBeforeUpdateOrderTest(unittest.TestCase):
+
+  def setUp(self):
+    os.environ["NEW_MODEL_DESIGN"] = "1"
+    os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+
+  def _engine(self, final_dropless: bool = False):
+    """The one-layer MoE engine with the dropless replay on; its loss raises the overflow flag on marked batches."""
+    cfg = _cfg(override_logical_axis_rules=True, ragged_buffer_factor=2.5)
+    mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    real_model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
+    with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=real_model):
+      engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh)
+    # The ring-of-experts prerequisites of retry_when_tokens_dropped cannot run on CPU, so the switches are set
+    # directly, as the dropless replay and final-dropless tests do.
+    engine._dropless_replay = True
+    engine._final_dropless = final_dropless
+    base_loss = maxtext_engine.make_router_replay_loss_fn(cfg)
+
+    def loss_fn(model, **batch):
+      program = _program(model)
+      loss, aux = base_loss(model, **batch)
+      aux = {k: v for k, v in aux.items() if k != "has_moe_overflow"}
+      marked = jnp.logical_or(batch["inputs"][0, 0] == _MARKED, batch["inputs"][0, 0] == _MARKED_B)
+      aux["has_moe_overflow"] = jnp.logical_and(marked, program == "regular")
+      aux["xent_sum"] = aux["xent_sum"] * _SCALES[program]
+      return loss, aux
+
+    engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn)
+    engine.with_loss_fn(loss_fn, has_aux=True)
+    engine.compile(None)
+    return engine
+
+  def _train(self, engine, staged: bool, prepare=None):
+    """Runs `_STEPS`; with `staged`, each next step's first micro-batch is prepared before the update.
+
+    Returns the final parameters, per-step metrics, and per staged boundary whether the last micro-batch's flag was
+    still pending when the next one was prepared.
+    """
+    prepare = prepare or engine.prepare_fwd_bwd_input
+    prepared = None
+    pending_at_prepare = []
+    metrics = []
+    for step, tokens in enumerate(_STEPS):
+      for j, first_token in enumerate(tokens):
+        final = j == len(tokens) - 1
+        if j == 0 and prepared is not None:
+          self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+          engine.fwd_bwd(prepared, apply_optimizer=False)
+          prepared = None
+        else:
+          engine.fwd_bwd(_payload(first_token), apply_optimizer=final)
+      if staged and step + 1 < len(_STEPS):
+        pending_at_prepare.append(engine._pending_overflow_check is not None)
+        prepared = prepare(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+        # Waiting resolves nothing: the flag is read, and the rerun run, by `update()`.
+        self.assertIsNotNone(engine._pending_overflow_check)
+      with mock.patch.object(engine, "record_metrics", wraps=engine.record_metrics) as record:
+        engine.update()
+      recorded = {call.args[0]: call.args[1] for call in record.call_args_list}
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append(
+          {
+              "replays": int(recorded["moe_dropless_replays"]),
+              "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+              "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+          }
+      )
+    self.assertIsNone(prepared)
+    params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+    return params, metrics, pending_at_prepare
+
+  def _assert_same_training(self, ref, new):
+    self.assertEqual(ref[1], new[1])
+    self.assertEqual(len(ref[0]), len(new[0]))
+    for r, n in zip(ref[0], new[0]):
+      np.testing.assert_array_equal(r, n)
+
+  def test_staged_before_a_replayed_update_trains_bit_identically(self):
+    reference = self._train(self._engine(), staged=False)
+    # Not vacuous: every step's last micro-batch replays (in update), and step 1's first (the prepared one) replays.
+    self.assertEqual([m["replays"] for m in reference[1]], [1, 2, 1])
+    engine = self._engine()
+    with mock.patch.object(engine, "_compile_for_batch", wraps=engine._compile_for_batch) as compiles:
+      staged = self._train(engine, staged=True)
+    self._assert_same_training(reference, staged)
+    # The last micro-batch's flag was still pending at both prepares, i.e. its rerun ran after the prepare.
+    self.assertEqual(staged[2], [True, True])
+    self.assertEqual(compiles.call_count, 1)  # no kernel was rebuilt for the prepared input
+
+  def test_with_the_final_dropless_program_too(self):
+    # retry_dropless_final_micro_batch (#5673): the last micro-batch runs the final program and cannot overflow, but
+    # its check is still created and read by update(), so the order is the same; only step 1's prepared first
+    # micro-batch replays.
+    reference = self._train(self._engine(final_dropless=True), staged=False)
+    self.assertEqual([m["replays"] for m in reference[1]], [0, 1, 0])
+    staged = self._train(self._engine(final_dropless=True), staged=True)
+    self._assert_same_training(reference, staged)
+    self.assertEqual(staged[2], [True, True])
+
+  def test_through_tunix_trainer_worker_the_deferred_update_trains_bit_identically(self):
+    """The same order, driven by a real Tunix `TrainerWorker`: `fwd_bwd(last, apply_optimizer=True,
+    defer_update=True)` runs the last micro-batch and holds the update, `stage_fwd_bwd` prepares the next step's first
+    micro-batch without waiting, `apply_deferred_update` waits for the device and runs the update (and with it the
+    last micro-batch's rerun), and `fwd_bwd_staged` trains the prepared micro-batch. Against the fused calls
+    `fwd_bwd(last, apply_optimizer=True)` and `fwd_bwd(next payload)`."""
+    from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
+    from tunix.experimental.worker import trainer_worker  # pylint: disable=import-outside-toplevel
+
+    if not hasattr(trainer_worker.TrainerWorker, "apply_deferred_update"):
+      self.skipTest("this Tunix has no TrainerWorker.apply_deferred_update (staging before the update)")
+
+    def run(staged: bool):
+      engine = self._engine()
+      worker = trainer_worker.TrainerWorker(trainer_factory=lambda: engine)
+      worker.initialize()
+      metrics, device_waits = [], []
+      staged_id = None
+      with mock.patch.object(engine, "fwd_bwd", wraps=engine.fwd_bwd) as fwd_bwd:
+        for step, tokens in enumerate(_STEPS):
+          for j, first_token in enumerate(tokens):
+            final = j == len(tokens) - 1
+            kwargs = {"apply_optimizer": True} if final else {}
+            if final and staged and step + 1 < len(_STEPS):
+              kwargs["defer_update"] = True
+            if j == 0 and staged_id is not None:
+              resp = worker.fwd_bwd_staged(staged_id, **kwargs)
+              staged_id = None
+            else:
+              request = datatypes.TrainRequest(request_id=f"s{step}m{j}", payload=_payload(first_token))
+              resp = worker.fwd_bwd(request, **kwargs)
+          if resp.metadata.get("update_deferred"):
+            nxt = datatypes.TrainRequest(request_id=f"s{step + 1}m0", payload=_payload(_STEPS[step + 1][0]))
+            stage = worker.stage_fwd_bwd(nxt)
+            self.assertTrue(stage.metadata["before_update"])
+            self.assertIsNone(stage.metadata["device_wait_s"])
+            update = worker.apply_deferred_update()
+            device_waits.append(update.metadata["device_wait_s"])
+            self.assertEqual(update.metadata["train_step"], step + 1)
+            staged_id = nxt.request_id
+          else:
+            self.assertEqual(resp.metadata["train_step"], step + 1)
+          buffer = engine.get_metrics(clear_cache=True)
+          metrics.append(
+              {
+                  "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+                  "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+              }
+          )
+        # `defer_update` is the worker's, never the engine's.
+        self.assertTrue(all("defer_update" not in c.kwargs for c in fwd_bwd.call_args_list))
+      self.assertIsNone(staged_id)
+      self.assertIsNone(worker._deferred_update)
+      params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+      return params, metrics, device_waits
+
+    reference = run(staged=False)
+    staged = run(staged=True)
+    self.assertEqual(reference[1], staged[1])
+    for r, n in zip(reference[0], staged[0]):
+      np.testing.assert_array_equal(r, n)
+    self.assertEqual(len(staged[2]), len(_STEPS) - 1)
+    self.assertTrue(all(w >= 0.0 for w in staged[2]))
+
+  def test_the_comparison_sees_a_rerun_of_the_wrong_micro_batch(self):
+    # Control: a rerun inside update() that ran on the prepared micro-batch instead of its own is caught.
+    reference = self._train(self._engine(), staged=False)
+    engine = self._engine()
+    holder = {}
+    real_prepare = engine.prepare_fwd_bwd_input
+    real_resolve = engine._resolve_overflow_check
+
+    def prepare(payload):
+      holder["prepared"] = real_prepare(payload)
+      return holder["prepared"]
+
+    def resolve():
+      check = engine._pending_overflow_check
+      prepared = holder.pop("prepared", None)
+      if check is not None and prepared is not None and check.micro_step > 0:
+        dynamic, _ = maxtext_engine._split_static_and_dynamic(prepared.batch)
+        engine._pending_overflow_check = dataclasses.replace(check, inputs=dynamic)
+      return real_resolve()
+
+    with mock.patch.object(engine, "_resolve_overflow_check", side_effect=resolve):
+      wrong = self._train(engine, staged=True, prepare=prepare)
+    self.assertNotEqual(reference[1], wrong[1])
+
+  def test_the_comparison_sees_a_different_prepared_micro_batch(self):
+    # Control: preparing the wrong micro-batch (the next step's second instead of its first) is caught.
+    reference = self._train(self._engine(), staged=False)
+    engine = self._engine()
+    real_prepare = engine.prepare_fwd_bwd_input
+    wrong = self._train(engine, staged=True, prepare=lambda payload: real_prepare(_payload(20)))
+    self.assertNotEqual(reference[1], wrong[1])
+
+
+if __name__ == "__main__":
+  unittest.main()