diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py
index 1bf9c5ae938..6959363509b 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -186,10 +186,11 @@ def sample(
     ) -> Optional[dict[str, Any]]:
         """Sample per-prompt trajectory groups intended for the current training step.
 
-        Only returns trajectories with target_weight_version == current_weight_version.
-        If insufficient trajectories are available, returns None to stall training
-        until the remaining trajectories are generated. This ensures no trajectory
-        loses its last chance to be used for its intended training step.
+        Only returns trajectories with target_weight_version >= current_weight_version,
+        so trajectories staged for a future step may be consumed early instead of
+        stalling. If insufficient trajectories are available, returns None to stall
+        training until the remaining trajectories are generated. This ensures no
+        trajectory loses its last chance to be used for its intended training step.
 
         Returns:
             Dictionary with 'trajectories' and 'avg_trajectory_age' keys, or None if insufficient data
@@ -252,7 +253,7 @@ def sample(
             intended_indices = [
                 i
                 for i in valid_indices
-                if self.target_weight_versions[i] == current_weight_version
+                if self.target_weight_versions[i] >= current_weight_version
             ]
 
             print(
@@ -283,8 +284,10 @@ def sample(
                 f"✅ Selected counts by generation weight-version: {Counter(sampled_weights)}"
             )
             print(f"📊 Average trajectory age: {avg_trajectory_age:.2f} steps")
+            sampled_targets = Counter(self.target_weight_versions[i] for i in selected)
             print(
-                f"🎯 All selected trajectories target step {current_weight_version} (100% target match)"
+                f"🎯 Selected trajectory targets for step {current_weight_version}: "
+                f"{dict(sampled_targets)}"
             )
 
             # Remove selected items in reverse order to maintain correct indices
diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py
index c0ce4446ec4..a68294cc2c4 100644
--- a/nemo_rl/algorithms/async_utils/trajectory_collector.py
+++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py
@@ -281,43 +281,51 @@ def _collection_loop(self):
                 if not self.running:
                     break
 
-                # Check if manually paused and wait
-                if not self._manual_pause_cleared.is_set() and self.running:
-                    self._manual_pause_cleared.wait()
-
-                # Check if refit is in progress and wait
-                if not self._refit_pause_cleared.is_set() and self.running:
-                    print("⏸️ Pausing collection for refit...")
-                    with self._efficiency_timer.time("idle/refit_event_wait"):
-                        self._refit_pause_cleared.wait()
-                    print("▶️ Refit completed, resuming collection")
-
-                # Check if generation limits require pausing collection
-                if self._should_pause_for_generation_limits() and self.running:
-                    # Only log warning once per weight version
-                    if self._last_limit_warning_version != self.current_weight_version:
-                        async_cfg = self.master_config.grpo.get("async_grpo", {})
-                        max_trajectory_age = async_cfg["max_trajectory_age_steps"]
-                        target_weights = [
-                            self.current_weight_version + i
-                            for i in range(max_trajectory_age)
-                        ]
-
-                        print(
-                            f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
-                            f"already exist in buffer. Waiting for weight update..."
-                        )
-                        self._last_limit_warning_version = self.current_weight_version
+                # Wait until NO pause condition holds, re-checking every
+                # condition after each wake-up: sequential checks race with
+                # pause() at validation boundaries and let full batches launch
+                # into the validation window.
+                while self.running:
+                    if not self._manual_pause_cleared.is_set():
+                        self._manual_pause_cleared.wait()
+                        continue  # re-check all conditions after waking
+
+                    if not self._refit_pause_cleared.is_set():
+                        print("⏸️ Pausing collection for refit...")
+                        with self._efficiency_timer.time("idle/refit_event_wait"):
+                            self._refit_pause_cleared.wait()
+                        print("▶️ Refit completed, resuming collection")
+                        continue  # re-check all conditions after waking
+
+                    if self._should_pause_for_generation_limits():
+                        # Only log warning once per weight version
+                        if (
+                            self._last_limit_warning_version
+                            != self.current_weight_version
+                        ):
+                            async_cfg = self.master_config.grpo.get("async_grpo", {})
+                            max_trajectory_age = async_cfg["max_trajectory_age_steps"]
+                            target_weights = [
+                                self.current_weight_version + i
+                                for i in range(max_trajectory_age)
+                            ]
+
+                            print(
+                                f"⏸️ Pausing collection: all target weights {target_weights} for weight version {self.current_weight_version} "
+                                f"already exist in buffer. Waiting for weight update..."
+                            )
+                            self._last_limit_warning_version = (
+                                self.current_weight_version
+                            )
 
-                        self._generation_limit_cleared.clear()  # Clear the event to pause
+                            self._generation_limit_cleared.clear()  # Clear the event to pause
 
-                    # Efficiently wait for generation limits to be cleared (no polling!)
-                    with self._efficiency_timer.time("idle/generation_limit_pause"):
-                        self._generation_limit_cleared.wait()
+                        # Efficiently wait for generation limits to be cleared (no polling!)
+                        with self._efficiency_timer.time("idle/generation_limit_pause"):
+                            self._generation_limit_cleared.wait()
+                        continue  # re-check all conditions after waking
 
-                    # Double-check we're still running after being woken up
-                    if not self.running:
-                        break
+                    break  # nothing requires pausing; clear to launch
 
                 if not self.running:
                     break
@@ -416,6 +424,12 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None:
 
             use_nemo_gym = _should_use_nemo_gym(self.master_config)
 
+            # Honor a manual pause (e.g. a validation boundary) before
+            # spawning, so the batch cannot launch into the val window.
+            if not self._manual_pause_cleared.is_set() and self.running:
+                print("⏸️ Manual pause before batch spawn: holding launch")
+                self._manual_pause_cleared.wait()
+
             if not self._refit_pause_cleared.is_set() and self.running:
                 with self._threads_lock:
                     active_threads = len(self._inflight_threads)
@@ -550,6 +564,10 @@ def resume_after_refit(self) -> None:
         # Invalidate&recompute vLLM caches after the in-flight weight updates if
         # recompute_kv_cache_after_weight_updates is True (AREAL-style implementation).
         # Otherwise, keep using the stale KV caches (Magistral-style implementation).
+        # NOTE: for drained (non-in-flight) refits with prefix caching enabled,
+        # cache invalidation is handled unconditionally in
+        # refit_policy_generation (grpo.py), which also covers the
+        # pre-validation refit path that never reaches this method.
         async_cfg = self.master_config.grpo.get("async_grpo", {})
         if async_cfg.get("in_flight_weight_updates", False) and async_cfg.get(
             "recompute_kv_cache_after_weight_updates", False
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index 12baec843cd..982a762d35f 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -1688,6 +1688,25 @@ def _scale(reward_tensor: torch.Tensor) -> torch.Tensor:
     return repeated_batch
 
 
+def _stable_group_ids(prompt_ids_for_adv, num_generations_per_prompt):
+    """Stable per-prompt grouping key for GRPO advantage computation.
+
+    GRPO groups samples by prompt (torch.unique) to compute the leave-one-out baseline. The default
+    key is the rendered prompt token-ids, but for agentic gym rollouts each generation's first-turn
+    prompt tokenizes slightly differently (observed on Qwen3-Instruct + hermes: every generation
+    becomes its own singleton group -> leave-one-out baseline == reward -> advantage == 0 -> zero
+    gradient; Exp 26). The training batch is laid out as contiguous num_gen blocks per prompt
+    (async: BatchedDataDict.from_batches of per-prompt groups; sync: repeat_interleave), so the
+    correct, model-agnostic group id is positional: index // num_gen. Falls back to the original
+    token-id grouping if the batch is not an exact multiple of num_gen (e.g. dynamic sampling).
+    """
+    n = int(prompt_ids_for_adv.shape[0])
+    g = int(num_generations_per_prompt)
+    if g <= 0 or n % g != 0:
+        return prompt_ids_for_adv
+    return (torch.arange(n, device=prompt_ids_for_adv.device) // g).unsqueeze(1)
+
+
 def extract_initial_prompt_messages(
     message_logs: list,
     original_prompt_lengths: torch.Tensor,
@@ -2218,7 +2237,9 @@ def refit_policy_generation(
     """
     synchronizer = getattr(policy_generation, "weight_synchronizer", None)
     if synchronizer is not None:
-        return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
+        sync_metrics = synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
+        _invalidate_prefix_cache_after_refit(policy_generation)
+        return sync_metrics
 
     # Megatron generation backend needs explicit suspend/resume around refits.
     if isinstance(policy_generation, MegatronGeneration):
@@ -2321,12 +2342,36 @@ def refit_policy_generation(
     if colocated_inference or isinstance(policy_generation, MegatronGeneration):
         policy_generation.prepare_for_generation(tags=["kv_cache"])
 
+    _invalidate_prefix_cache_after_refit(policy_generation)
+
     if isinstance(policy_generation, MegatronGeneration):
         policy_generation.resume_after_refit()
 
     return {}
 
 
+def _invalidate_prefix_cache_after_refit(
+    policy_generation: GenerationInterface,
+) -> None:
+    """Drop reusable KV blocks after a weight update.
+
+    vLLM prefix-cache blocks are keyed only by token ids, so blocks
+    prefilled under the old weights would be silently reused after refit
+    (stale KV -> train/gen logprob divergence). Called on every refit path;
+    backends without reusable caches inherit the no-op interface default.
+    """
+    generation_cfg = getattr(policy_generation, "cfg", None) or {}
+    vllm_cfg = generation_cfg.get("vllm_cfg") or {}
+    if vllm_cfg.get("enable_prefix_caching"):
+        if not policy_generation.invalidate_kv_cache():
+            raise RuntimeError(
+                "❌ Error: prefix caching is enabled but invalidating the "
+                "vLLM prefix/KV cache after refit failed; continuing would "
+                "sample rollouts against stale KV computed under the "
+                "pre-refit weights."
+            )
+
+
 def _initial_policy_generation_stale(
     policy_generation: GenerationInterface, completed_steps: int
 ) -> bool:
@@ -2450,6 +2495,30 @@ def compute_and_apply_seq_logprob_error_masking(
             )
             masked_correct_pct = masked_correct_count / num_masked_seqs
 
+            # [lp-mask-debug] one parseable line per step attributing train/gen
+            # logprob divergence; opt in via NRL_LP_MASK_DEBUG=1.
+            if os.environ.get("NRL_LP_MASK_DEBUG") == "1":
+                seq_lens = mask.sum(dim=-1)
+                masked_rows = [
+                    (
+                        int(seq_lens[i]),
+                        float(seq_mult_prob_error[i]),
+                        float(rewards.view(-1)[i]),
+                    )
+                    for i in torch.nonzero(diff_mask_bool).flatten().tolist()
+                ]
+                kept_bool = seq_error_mask.bool() & valid_seq_mask
+                kept_lens = seq_lens[kept_bool]
+                print(
+                    "[lp-mask-debug] masked(len,err,rew)="
+                    + ";".join(f"{l},{e:.2f},{r:.0f}" for l, e, r in masked_rows[:200])
+                    + f" | kept_len mean={float(kept_lens.float().mean()):.0f}"
+                    f" p90={float(kept_lens.float().quantile(0.9)):.0f}"
+                    f" max={int(kept_lens.max())}"
+                    f" | masked_len mean={sum(r[0] for r in masked_rows) / len(masked_rows):.0f}",
+                    flush=True,
+                )
+
         # Compute after-mask metrics (only for sequences that passed the threshold)
         kept_mask = seq_error_mask.bool() & valid_seq_mask
         if kept_mask.sum() > 0:
@@ -3064,8 +3133,17 @@ def grpo_train(
                     sample_mask = train_data["sample_mask"]
                     mask = token_mask * sample_mask.unsqueeze(-1)
 
+                    # Positional grouping is only needed for agentic gym rollouts,
+                    # where per-generation prompt tokenization is non-deterministic;
+                    # keep main's token-id grouping for every other GRPO user.
+                    advantage_group_ids = prompt_ids_for_adv
+                    if _should_use_nemo_gym(master_config):
+                        advantage_group_ids = _stable_group_ids(
+                            prompt_ids_for_adv,
+                            master_config.grpo["num_generations_per_prompt"],
+                        )
                     train_data["advantages"] = adv_estimator.compute_advantage(
-                        prompt_ids=prompt_ids_for_adv,
+                        prompt_ids=advantage_group_ids,
                         rewards=rewards,
                         mask=mask,
                         repeated_batch=repeated_batch,
@@ -3972,14 +4050,6 @@ def async_grpo_train(
         next_nemo_gym_task_index=next_nemo_gym_task_index,
     )
 
-    # Start trajectory collection in background
-    collection_task = trajectory_collector.start_collection.remote(dataloader)
-
-    # Ensure collector knows initial weight version
-    trajectory_collector.set_weight_version.remote(weight_version)
-
-    print("📦 Started continuous background trajectory collection")
-
     print(
         f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, max_age={max_trajectory_age_steps} steps"
     )
@@ -4009,6 +4079,17 @@ def async_grpo_train(
             traceback.print_exc()
             return
 
+    # Start trajectory collection only after generation holds real weights.
+    # The engines come up with load_format=dummy (weights arrive via the refit
+    # above); collecting before the refit fills the buffer with garbage
+    # rollouts sampled from randomly initialized weights.
+    collection_task = trajectory_collector.start_collection.remote(dataloader)  # noqa: F841
+
+    # Ensure collector knows initial weight version
+    trajectory_collector.set_weight_version.remote(weight_version)
+
+    print("📦 Started continuous background trajectory collection")
+
     print("✅ Policy generation setup complete, proceeding to validation...")
 
     # Run validation at start if configured
@@ -4448,8 +4529,17 @@ def async_grpo_train(
                     sample_mask = train_data["sample_mask"]
                     mask = token_mask * sample_mask.unsqueeze(-1)
 
+                    # Positional grouping is only needed for agentic gym rollouts,
+                    # where per-generation prompt tokenization is non-deterministic;
+                    # keep main's token-id grouping for every other GRPO user.
+                    advantage_group_ids = prompt_ids_for_adv
+                    if _should_use_nemo_gym(master_config):
+                        advantage_group_ids = _stable_group_ids(
+                            prompt_ids_for_adv,
+                            master_config.grpo["num_generations_per_prompt"],
+                        )
                     train_data["advantages"] = adv_estimator.compute_advantage(
-                        prompt_ids=prompt_ids_for_adv,
+                        prompt_ids=advantage_group_ids,
                         rewards=rewards,
                         mask=mask,
                         repeated_batch=repeated_batch,
@@ -4558,6 +4648,12 @@ def async_grpo_train(
                     with timer.time("idle/validation"):
                         # Pause trajectory collection during validation to reduce memory pressure
                         trajectory_collector.pause.remote()
+                        # Drain in-flight rollouts too: pause only stops new
+                        # launches, and in-flight train rollouts sharing the
+                        # engines push val agents into timeout.
+                        ray.get(
+                            trajectory_collector.wait_for_pending_generations.remote()
+                        )
 
                         if NEED_REFIT and POLICY_GENERATION_STALE:
                             refit_metrics = refit_policy_generation(
diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py
index 9a17d294dd9..9adc061f047 100644
--- a/nemo_rl/environments/nemo_gym.py
+++ b/nemo_rl/environments/nemo_gym.py
@@ -11,15 +11,18 @@
 # 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.
+import asyncio
 import math
 import os
 import subprocess
 import sys
 from collections import Counter
 from collections.abc import AsyncGenerator
+
 from pathlib import Path
 from typing import Any, Dict, List, NotRequired, Optional, TypedDict
 
+import aiohttp
 import ray
 import torch
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
@@ -111,6 +114,12 @@ class NemoGymConfig(TypedDict):
     # Forwarded from policy.tokenizer.use_fastokens so rollout actors patch their
     # tokenizer consistently with the driver. Defaults to off when absent.
     use_fastokens: NotRequired[bool]
+    # When true, rollouts that fail (transport errors, NaN generation
+    # logprobs, malformed or empty generation data) are back-filled as
+    # zero-reward trajectories so the batch keeps its shape. Absent/false
+    # keeps the default fail-fast behavior (the whole batch raises); long
+    # multi-node benchmark runs enable it.
+    backfill_failed_rollouts: NotRequired[bool]
 
 
 def _detect_invalid_tool_call_and_malformed_thinking(
@@ -291,17 +300,47 @@ async def run_rollouts(
 
         timer = Timer()
         counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples)
+        backfill_failed_rollouts = bool(self.cfg.get("backfill_failed_rollouts"))
 
         timer.start("_run_rollouts_total")
         nemo_gym_result_iterator = self.rch.run_examples(
             examples=nemo_gym_examples, head_server_config=self.head_server_config
         )
 
+        def _final_timing_metrics() -> dict:
+            timer.stop("_run_rollouts_total")
+            final_metrics = timer.get_timing_metrics("sum")
+            total_time = final_metrics.pop("_run_rollouts_total")
+            # The postprocess timer label never fires when every rollout failed
+            # at `await task` and was back-filled, so default it to 0 rather
+            # than crash.
+            final_metrics[f"{timer_prefix}/postprocess_results_pct"] = (
+                100
+                * final_metrics.get(f"{timer_prefix}/postprocess_results", 0.0)
+                / total_time
+                if total_time
+                else 0.0
+            )
+            return final_metrics
+
         num_results = 0
+        seen_rowidxs: set[int] = set()
         for task in nemo_gym_result_iterator:
             with timer.time(label=f"{timer_prefix}/await_results"):
                 try:
                     nemo_gym_row, nemo_gym_result = await task
+                except (aiohttp.ClientError, asyncio.TimeoutError) as error:
+                    if not backfill_failed_rollouts:
+                        raise
+                    # Back-fill the missing row after the stream instead of
+                    # aborting the batch.
+                    print(
+                        f"  [nemo_gym] WARNING: rollout failed "
+                        f"({type(error).__name__}: {error}); will back-fill as "
+                        "a zero-reward trajectory instead of aborting the batch.",
+                        flush=True,
+                    )
+                    continue
                 except Exception as error:
                     if hasattr(error, "response_content"):
                         print(
@@ -312,23 +351,46 @@ async def run_rollouts(
                     raise
 
             with timer.time(label=f"{timer_prefix}/postprocess_results"):
-                nemo_rl_result = self._postprocess_nemo_gym_to_nemo_rl_result(
-                    nemo_gym_result, tokenizer
-                )
+                try:
+                    nemo_rl_result = self._postprocess_nemo_gym_to_nemo_rl_result(
+                        nemo_gym_result, tokenizer
+                    )
+                except Exception as error:
+                    if not backfill_failed_rollouts:
+                        raise
+                    print(
+                        f"  [nemo_gym] WARNING: failed to postprocess rollout "
+                        f"{nemo_gym_row.get('_rowidx', '<unknown>')} "
+                        f"({type(error).__name__}: {error}); "
+                        "back-filling as a zero-reward trajectory.",
+                        flush=True,
+                    )
+                    nemo_rl_result = self._zero_reward_nemo_rl_result(
+                        tokenizer,
+                        nemo_gym_result if isinstance(nemo_gym_result, dict) else None,
+                        reason=f"postprocess_failed:{type(error).__name__}",
+                    )
                 if _has_nan_generation_logprobs(nemo_rl_result):
-                    raise RuntimeError("Generation logprobs contain NaN")
+                    if not backfill_failed_rollouts:
+                        raise RuntimeError("Generation logprobs contain NaN")
+                    print(
+                        f"  [nemo_gym] WARNING: rollout "
+                        f"{nemo_gym_row.get('_rowidx', '<unknown>')} returned NaN "
+                        "generation logprobs; back-filling as a zero-reward "
+                        "trajectory instead of aborting the batch.",
+                        flush=True,
+                    )
+                    nemo_rl_result = self._zero_reward_nemo_rl_result(
+                        tokenizer,
+                        nemo_gym_result if isinstance(nemo_gym_result, dict) else None,
+                        reason="nan_generation_logprobs",
+                    )
 
             num_results += 1
+            seen_rowidxs.add(nemo_gym_row["_rowidx"])
             timing_metrics = None
             if num_results == len(nemo_gym_examples):
-                timer.stop("_run_rollouts_total")
-                timing_metrics = timer.get_timing_metrics("sum")
-                total_time = timing_metrics.pop("_run_rollouts_total")
-                timing_metrics[f"{timer_prefix}/postprocess_results_pct"] = (
-                    100
-                    * timing_metrics[f"{timer_prefix}/postprocess_results"]
-                    / total_time
-                )
+                timing_metrics = _final_timing_metrics()
 
             agent_name = nemo_gym_row["agent_ref"]["name"]
             counts_left[agent_name] -= 1
@@ -348,6 +410,80 @@ async def run_rollouts(
 
             yield nemo_gym_row["_rowidx"], nemo_rl_result, timing_metrics
 
+        # Preserve batch shape when rollouts were dropped at `await task`: the
+        # stream consumer expects exactly one result per input row, so emit the
+        # missing rows as zero-reward trajectories.
+        if not backfill_failed_rollouts:
+            return
+        missing_rowidxs = [
+            rowidx
+            for rowidx in range(len(nemo_gym_examples))
+            if rowidx not in seen_rowidxs
+        ]
+        if missing_rowidxs:
+            print(
+                f"  [nemo_gym] WARNING: back-filling "
+                f"{len(missing_rowidxs)}/{len(nemo_gym_examples)} failed "
+                "rollout(s) as zero-reward trajectories (batch counts preserved).",
+                flush=True,
+            )
+        for rowidx in missing_rowidxs:
+            num_results += 1
+            timing_metrics = (
+                _final_timing_metrics()
+                if num_results == len(nemo_gym_examples)
+                else None
+            )
+            yield (
+                rowidx,
+                self._zero_reward_nemo_rl_result(tokenizer, reason="rollout_failed"),
+                timing_metrics,
+            )
+
+    def _zero_reward_nemo_rl_result(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        nemo_gym_result: dict | None = None,
+        reason: str | None = None,
+    ) -> dict:
+        fallback_token = (
+            tokenizer.pad_token_id
+            if tokenizer.pad_token_id is not None
+            else (tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0)
+        )
+        if nemo_gym_result is None:
+            nemo_gym_result = {}
+
+        response = nemo_gym_result.get("response")
+        if not isinstance(response, dict):
+            response = {}
+            nemo_gym_result["response"] = response
+        if not isinstance(response.get("output"), list):
+            response["output"] = []
+        nemo_gym_result["reward"] = 0.0
+        if reason:
+            nemo_gym_result["nemo_rl_fallback_reason"] = reason
+
+        message_log = [
+            {
+                "role": "user",
+                "content": "",
+                "token_ids": torch.tensor(
+                    [fallback_token, fallback_token], dtype=torch.int64
+                ),
+            },
+            {
+                "role": "assistant",
+                "content": "",
+                "token_ids": torch.tensor([fallback_token], dtype=torch.int64),
+            },
+        ]
+        return {
+            "message_log": message_log,
+            "input_message_log": message_log[:1],
+            "full_result": nemo_gym_result,
+        }
+
     def _postprocess_nemo_gym_to_nemo_rl_result(
         self, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase
     ) -> dict:
@@ -355,15 +491,31 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
             f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}"
         )
 
+        response = nemo_gym_result.get("response")
+        response_output = response.get("output") if isinstance(response, dict) else None
+        if not isinstance(response_output, list):
+            if not self.cfg.get("backfill_failed_rollouts"):
+                raise ValueError(
+                    "NeMo Gym returned a malformed response.output: "
+                    f"{type(response_output).__name__}"
+                )
+            return self._zero_reward_nemo_rl_result(
+                tokenizer,
+                nemo_gym_result,
+                reason=f"malformed_response_output:{type(response_output).__name__}",
+            )
+
         nemo_rl_message_log = []
         seen_token_ids: List[int] = []
         batch_decode_items = []
-        for output_item_dict in nemo_gym_result["response"]["output"]:
+        for output_item_dict in response_output:
             # Nemo RL really only has two types of messages: assistant and not assistant since that is all that it is concerned with (i.e. to train or not to train)
             # Here we map all the trainable messages to assistant and all the non-trainable messages to user.
             # Eventually we can maybe be smarter about this, but this is functional for now.
 
             # Note that NeMo-Gym will only return token ids on "assistant" messages and not other message types.
+            if not isinstance(output_item_dict, dict):
+                continue
             # Also skip if generation_token_ids is present but empty, e.g. all-EOS generation stripped to [] — torch.tensor([]) defaults to float32 and breaks batch dtype consistency.
             if (
                 "generation_token_ids" not in output_item_dict
@@ -479,20 +631,26 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
                 output_item_dict["generation_str"] = generation_str
 
         if not nemo_rl_message_log:
-            input_messages = nemo_gym_result["responses_create_params"]["input"]
+            input_messages = nemo_gym_result.get("responses_create_params", {}).get(
+                "input"
+            )
             try:
-                prompt_token_ids = tokenizer.apply_chat_template(
-                    input_messages, tokenize=True
-                )
-                prompt_len_str = f"{len(prompt_token_ids)} tokens"
+                if input_messages is None:
+                    prompt_len_str = "<unknown>"
+                else:
+                    prompt_token_ids = tokenizer.apply_chat_template(
+                        input_messages, tokenize=True
+                    )
+                    prompt_len_str = f"{len(prompt_token_ids)} tokens"
             except Exception as e:
                 prompt_len_str = (
-                    f"<unknown — apply_chat_template failed: {type(e).__name__}: {e}>"
+                    f"<unknown - apply_chat_template failed: {type(e).__name__}: {e}>"
                 )
             output_item_types = [
-                o.get("type") for o in nemo_gym_result["response"]["output"]
+                o.get("type") if isinstance(o, dict) else type(o).__name__
+                for o in response_output
             ]
-            raise ValueError(
+            no_generation_message = (
                 f"NeMo Gym returned a result with no generation data. "
                 f"Possible causes: (1) the prompt for the first turn already exceeds the vLLM max_model_len, "
                 f"so vLLM rejected the request before any tokens could be generated; "
@@ -503,6 +661,16 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
                 f"above the prompt length above.\n"
                 f"  → If (2): inspect why no assistant content was produced for this rollout."
             )
+            if not self.cfg.get("backfill_failed_rollouts"):
+                raise ValueError(no_generation_message)
+            print(
+                no_generation_message
+                + "\n  Treating this rollout as a masked zero-reward trajectory instead of aborting the batch.",
+                flush=True,
+            )
+            return self._zero_reward_nemo_rl_result(
+                tokenizer, nemo_gym_result, reason="no_generation_data"
+            )
 
         return {
             "message_log": nemo_rl_message_log,
diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py
index 679cb37752d..b4dd4fc908e 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker_async.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py
@@ -468,7 +468,12 @@ async def preprocess_chat(
                     if message.get("tool_calls"):
                         message["tool_calls"] = list(message["tool_calls"])
 
-                messages_for_replace_prefix_tokens = deepcopy(messages)
+                    content = message.get("content")
+                    if content is not None and not isinstance(content, (list, str)):
+                        try:
+                            message["content"] = list(content)
+                        except TypeError:
+                            message["content"] = []
 
                 # Temporarily set to 1 so vLLM's pre-tokenization length check passes;
                 # the actual value will be set through _clamp_max_tokens later.
@@ -518,6 +523,27 @@ async def preprocess_chat(
                         )
                     return res
 
+                # vLLM normalizes reasoning and tool-call message content during
+                # preprocessing.  Reuse that representation for the isolated
+                # prefix render, while removing NeMo-RL's token bookkeeping.
+                excluded_fields = {
+                    "prompt_token_ids",
+                    "generation_token_ids",
+                    "generation_log_probs",
+                }
+                messages_for_replace_prefix_tokens = []
+                for message in messages:
+                    if isinstance(message, dict):
+                        messages_for_replace_prefix_tokens.append(
+                            {
+                                key: deepcopy(value)
+                                for key, value in message.items()
+                                if key not in excluded_fields
+                            }
+                        )
+                    else:
+                        messages_for_replace_prefix_tokens.append(deepcopy(message))
+
                 last_assistant_message_idx = None
                 for i in reversed(range(len(messages_for_replace_prefix_tokens))):
                     if messages_for_replace_prefix_tokens[i]["role"] == "assistant":
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index d5aaae2d6fc..00e6cb8b4b5 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -1167,7 +1167,9 @@ def _create_megatron_config(
         ),
         optimizer=OptimizerConfig(**optimizer_kwargs),
         ddp=DistributedDataParallelConfig(
-            check_for_nan_in_grad=True,
+            check_for_nan_in_grad=config["megatron_cfg"][
+                "distributed_data_parallel_config"
+            ].get("check_for_nan_in_grad", True),
             grad_reduce_in_fp32=config["megatron_cfg"][
                 "distributed_data_parallel_config"
             ]["grad_reduce_in_fp32"],
diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py
index a84187ef0a1..695ba2654f6 100644
--- a/nemo_rl/models/policy/__init__.py
+++ b/nemo_rl/models/policy/__init__.py
@@ -281,6 +281,7 @@ class MegatronDDPConfig(TypedDict):
     overlap_param_gather: bool
     use_custom_fsdp: bool
     data_parallel_sharding_strategy: str
+    check_for_nan_in_grad: NotRequired[bool]
 
 
 class Fp8Config(TypedDict):
diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py
index b889dc42248..28dfc2b8861 100644
--- a/tests/unit/algorithms/test_async_utils.py
+++ b/tests/unit/algorithms/test_async_utils.py
@@ -644,6 +644,31 @@ def test_replay_buffer_target_weight_matching(self):
 
         ray.kill(buffer)
 
+    def test_replay_buffer_future_target_weight_matching(self):
+        """Test that sampling can use trajectories intended for a future step."""
+        buffer = ReplayBuffer.remote(max_size=10)
+
+        trajectory = {
+            "batch": {"data": "for_future_step"},
+            "rollout_metrics": {"reward": 1.0},
+        }
+        ray.get(
+            buffer.add.remote(trajectory, weight_version=1, target_weight_version=3)
+        )
+
+        sample_result = ray.get(
+            buffer.sample.remote(
+                num_prompt_groups=1,
+                current_weight_version=2,
+                max_age_steps=1,
+            )
+        )
+
+        assert sample_result is not None
+        assert sample_result["trajectories"][0]["batch"]["data"] == "for_future_step"
+
+        ray.kill(buffer)
+
     def test_replay_buffer_get_existing_target_weights(self):
         """Test getting existing target weight versions."""
         buffer = ReplayBuffer.remote(max_size=10)
diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py
index cc564c6daa5..f6f5e50e3fc 100644
--- a/tests/unit/algorithms/test_grpo.py
+++ b/tests/unit/algorithms/test_grpo.py
@@ -37,10 +37,13 @@
     _raise_if_reward_penalties_enabled_without_nemo_gym,
     _resolve_message_level_advantage_penalties,
     _should_use_async_rollouts,
+    _stable_group_ids,
+    add_grpo_token_loss_masks_and_generation_logprobs,
     aggregate_rollout_metrics,
     async_grpo_train,
     compute_and_apply_seq_logprob_error_masking,
     dynamic_sampling,
+    extract_initial_prompt_messages,
     grpo_train,
     refit_policy_generation,
     validate,
@@ -156,6 +159,29 @@ def test_initial_policy_generation_stale() -> None:
     assert _initial_policy_generation_stale(generation, completed_steps=0)
 
 
+def test_stable_group_ids_uses_contiguous_prompt_groups():
+    rendered_prompt_ids = torch.tensor(
+        [
+            [10, 11],
+            [10, 12],
+            [20, 21],
+            [20, 22],
+        ]
+    )
+
+    result = _stable_group_ids(rendered_prompt_ids, num_generations_per_prompt=2)
+
+    assert torch.equal(result, torch.tensor([[0], [0], [1], [1]]))
+
+
+def test_stable_group_ids_falls_back_for_incomplete_group():
+    rendered_prompt_ids = torch.tensor([[10], [11], [20]])
+
+    result = _stable_group_ids(rendered_prompt_ids, num_generations_per_prompt=2)
+
+    assert result is rendered_prompt_ids
+
+
 @pytest.fixture
 def mock_grpo_components():
     # Create mock components
@@ -2442,6 +2468,13 @@ def test_grpo_train_skips_prev_logprobs_when_force_on_policy_ratio(
     )
 
 
+@pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train])
+@pytest.mark.parametrize(
+    ("val_at_end", "expected_validation_steps"),
+    [(False, [4]), (True, [4, 5])],
+)
+
+
 @pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train])
 def test_grpo_exit_on_max_steps(mock_grpo_components, train_func):
     """Test that GRPO training loop exits when max_num_steps is reached"""
@@ -3629,3 +3662,138 @@ def test_mixed_metrics(self):
         assert result["total_turns"] == 45
         assert result["accuracy"] == pytest.approx(0.8)
         assert result["min_accuracy_rate"] == pytest.approx(0.2)
+
+
+class TestMultiTurnPromptHelpers:
+    """Tests for GRPO multi-turn prompt extraction and loss masking."""
+
+    def test_prompt_extraction_with_multi_turn_history(self):
+        original_prompt_messages = [
+            {
+                "role": "user",
+                "content": "What is 2+2?",
+                "token_ids": torch.tensor([1, 2]),
+            },
+            {
+                "role": "assistant",
+                "content": "4",
+                "token_ids": torch.tensor([3, 4]),
+            },
+            {
+                "role": "user",
+                "content": "Now what is 3+3?",
+                "token_ids": torch.tensor([5, 6]),
+            },
+        ]
+        generated_message = {
+            "role": "assistant",
+            "content": "6",
+            "token_ids": torch.tensor([7, 8]),
+        }
+        full_message_log = original_prompt_messages + [generated_message]
+        original_prompt_length = sum(
+            len(message["token_ids"]) for message in original_prompt_messages
+        )
+
+        result = extract_initial_prompt_messages(
+            [full_message_log], torch.tensor([original_prompt_length])
+        )
+
+        assert [message["role"] for message in result[0]] == [
+            "user",
+            "assistant",
+            "user",
+        ]
+        assert generated_message not in result[0]
+
+    def test_prompt_extraction_with_system_message(self):
+        original_prompt_messages = [
+            {
+                "role": "system",
+                "content": "You are a math tutor.",
+                "token_ids": torch.tensor([1, 2, 3]),
+            },
+            {
+                "role": "user",
+                "content": "What is 2+2?",
+                "token_ids": torch.tensor([4, 5]),
+            },
+        ]
+        generated_message = {
+            "role": "assistant",
+            "content": "4",
+            "token_ids": torch.tensor([6, 7]),
+        }
+        full_message_log = original_prompt_messages + [generated_message]
+        original_prompt_length = sum(
+            len(message["token_ids"]) for message in original_prompt_messages
+        )
+
+        result = extract_initial_prompt_messages(
+            [full_message_log], torch.tensor([original_prompt_length])
+        )
+
+        assert [message["role"] for message in result[0]] == ["system", "user"]
+        assert generated_message not in result[0]
+
+    def test_grpo_loss_mask_excludes_assistant_prompt_history(self):
+        message_log = [
+            {
+                "role": "user",
+                "content": "What is 2+2?",
+                "token_ids": torch.tensor([1, 2]),
+            },
+            {
+                "role": "assistant",
+                "content": "4",
+                "token_ids": torch.tensor([3, 4]),
+            },
+            {
+                "role": "user",
+                "content": "Now what is 3+3?",
+                "token_ids": torch.tensor([5, 6]),
+            },
+            {
+                "role": "assistant",
+                "content": "6",
+                "token_ids": torch.tensor([7, 8]),
+                "generation_logprobs": torch.tensor([0.1, 0.2]),
+            },
+        ]
+
+        add_grpo_token_loss_masks_and_generation_logprobs([message_log])
+
+        assert torch.equal(message_log[0]["token_loss_mask"], torch.tensor([0, 0]))
+        assert torch.equal(message_log[1]["token_loss_mask"], torch.tensor([0, 0]))
+        assert torch.equal(message_log[2]["token_loss_mask"], torch.tensor([0, 0]))
+        assert torch.equal(message_log[3]["token_loss_mask"], torch.tensor([1, 1]))
+
+    def test_grpo_loss_mask_uses_generation_logprobs_marker(self):
+        message_log = [
+            {
+                "role": "assistant",
+                "content": "prompt history",
+                "token_ids": torch.tensor([1, 2]),
+            },
+            {
+                "role": "user",
+                "content": "next question",
+                "token_ids": torch.tensor([3, 4]),
+                "generation_logprobs": torch.tensor([0.3, 0.4]),
+            },
+            {
+                "role": "assistant",
+                "content": "generated response",
+                "token_ids": torch.tensor([5, 6]),
+                "generation_logprobs": torch.tensor([0.5, 0.6]),
+            },
+        ]
+
+        add_grpo_token_loss_masks_and_generation_logprobs([message_log])
+
+        assert torch.equal(message_log[0]["token_loss_mask"], torch.tensor([0, 0]))
+        assert torch.equal(
+            message_log[0]["generation_logprobs"], torch.tensor([0.0, 0.0])
+        )
+        assert torch.equal(message_log[1]["token_loss_mask"], torch.tensor([0, 0]))
+        assert torch.equal(message_log[2]["token_loss_mask"], torch.tensor([1, 1]))
diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py
index a4ca5759e98..3e322a2fe98 100644
--- a/tests/unit/environments/test_nemo_gym.py
+++ b/tests/unit/environments/test_nemo_gym.py
@@ -11,15 +11,18 @@
 # 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.
+import asyncio
 import json
 import time
 from copy import deepcopy
 from pathlib import Path
+from types import MethodType
 
 import pytest
 import ray
 import requests
 import torch
+from aiohttp import ClientConnectionError
 from yaml import safe_load
 
 from nemo_rl.algorithms.grpo import MasterConfig
@@ -311,67 +314,192 @@ class _MockSelf:
     assert nemo_gym_result["response"]["output"][1]["generation_str"] == "6 7"
 
 
-def test_nemo_gym_postprocess_no_generation_data_raises():
-    """When no output item carries generation data, the postprocess should raise a
-    ValueError that reports the prompt length and the response.output item types."""
-
+def test_nemo_gym_postprocess_empty_generation_is_masked_zero_reward():
     class _Tokenizer:
-        def apply_chat_template(self, input_messages, tokenize=True):
-            # Pretend the prompt is 1234 tokens long.
-            return list(range(1234))
+        pad_token_id = 7
+        eos_token_id = 8
 
-    nemo_gym_result = {
-        "response": {
-            "output": [
-                {"type": "reasoning"},
-                {"type": "function_call"},
-            ]
-        },
-        "responses_create_params": {"input": [{"role": "user", "content": "hi"}]},
-    }
+        def apply_chat_template(self, messages, tokenize=True):
+            return [1, 2, 3]
 
     class _MockSelf:
-        cfg = {}
+        cfg = {"backfill_failed_rollouts": True}
 
-    with pytest.raises(ValueError) as excinfo:
+    nemo_gym_result = {
+        "response": {"output": []},
+        "responses_create_params": {"input": [{"role": "user", "content": "x"}]},
+        "reward": 1.0,
+    }
+
+    mock_self = _MockSelf()
+    zero_reward = NemoGym.__ray_metadata__.modified_class._zero_reward_nemo_rl_result
+    mock_self._zero_reward_nemo_rl_result = MethodType(zero_reward, mock_self)
+    result = (
         NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
-            _MockSelf(), nemo_gym_result, _Tokenizer()
+            mock_self, nemo_gym_result, _Tokenizer()
         )
+    )
+
+    assert result["full_result"]["reward"] == 0.0
+    assert result["message_log"][0]["token_ids"].tolist() == [7, 7]
+    assert result["message_log"][1]["token_ids"].tolist() == [7]
+    assert "generation_logprobs" not in result["message_log"][1]
+
 
-    msg = str(excinfo.value)
-    assert "no generation data" in msg
-    assert "1234 tokens" in msg
-    # The error surfaces the response.output item types to help diagnose case (2).
-    assert "['reasoning', 'function_call']" in msg
+def _collect_rollout_stream(mock_self, rows, tokenizer):
+    """Drain the run_rollouts async generator and return the yielded items."""
 
+    async def _collect():
+        items = []
+        async for item in NemoGym.__ray_metadata__.modified_class.run_rollouts(
+            mock_self, rows, tokenizer, "timing/rollout"
+        ):
+            items.append(item)
+        return items
 
-def test_nemo_gym_postprocess_no_generation_data_chat_template_failure():
-    """If apply_chat_template itself fails while building the error message, the
-    postprocess should still raise the original 'no generation data' ValueError with
-    the prompt length reported as unknown rather than masking it with a new error."""
+    return asyncio.run(_collect())
+
+
+def test_nemo_gym_transport_failure_preserves_batch_shape():
+    async def failed_rollout():
+        raise ClientConnectionError("transient failure")
+
+    class _RolloutCollectionHelper:
+        def run_examples(self, examples, head_server_config):
+            return [failed_rollout()]
 
     class _Tokenizer:
-        def apply_chat_template(self, input_messages, tokenize=True):
-            raise RuntimeError("boom")
+        pad_token_id = 7
+        eos_token_id = 8
 
-    nemo_gym_result = {
-        "response": {"output": [{"type": "reasoning"}]},
-        "responses_create_params": {"input": [{"role": "user", "content": "hi"}]},
-    }
+    class _MockSelf:
+        cfg = {"backfill_failed_rollouts": True}
+        rch = _RolloutCollectionHelper()
+        head_server_config = object()
+
+    mock_self = _MockSelf()
+    modified_class = NemoGym.__ray_metadata__.modified_class
+    mock_self._postprocess_nemo_gym_to_nemo_rl_result = MethodType(
+        modified_class._postprocess_nemo_gym_to_nemo_rl_result, mock_self
+    )
+    mock_self._zero_reward_nemo_rl_result = MethodType(
+        modified_class._zero_reward_nemo_rl_result, mock_self
+    )
+
+    items = _collect_rollout_stream(
+        mock_self, [{"_rowidx": 0, "agent_ref": {"name": "agent"}}], _Tokenizer()
+    )
+
+    assert len(items) == 1
+    rowidx, result, timing_metrics = items[0]
+    assert rowidx == 0
+    assert result["full_result"]["reward"] == 0.0
+    assert result["full_result"]["nemo_rl_fallback_reason"] == "rollout_failed"
+    # The back-filled final row still carries the batch timing metrics.
+    assert timing_metrics is not None
+
+
+def test_nemo_gym_transport_failure_raises_by_default():
+    """Without backfill_failed_rollouts, rollout failures fail the batch."""
+
+    async def failed_rollout():
+        raise ClientConnectionError("transient failure")
+
+    class _RolloutCollectionHelper:
+        def run_examples(self, examples, head_server_config):
+            return [failed_rollout()]
+
+    class _Tokenizer:
+        pad_token_id = 7
+        eos_token_id = 8
 
     class _MockSelf:
         cfg = {}
+        rch = _RolloutCollectionHelper()
+        head_server_config = object()
 
-    with pytest.raises(ValueError) as excinfo:
-        NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result(
-            _MockSelf(), nemo_gym_result, _Tokenizer()
+    with pytest.raises(ClientConnectionError):
+        _collect_rollout_stream(
+            _MockSelf(), [{"_rowidx": 0, "agent_ref": {"name": "agent"}}], _Tokenizer()
+        )
+
+
+def test_nemo_gym_postprocess_failure_preserves_batch_shape():
+    async def malformed_rollout():
+        return (
+            {"_rowidx": 0, "agent_ref": {"name": "agent"}},
+            {
+                "response": {
+                    "output": [
+                        {
+                            "generation_token_ids": [3],
+                            "generation_log_probs": [-0.1],
+                        }
+                    ]
+                },
+                "reward": 1.0,
+            },
         )
 
-    msg = str(excinfo.value)
-    assert "no generation data" in msg
-    assert "apply_chat_template failed" in msg
-    assert "RuntimeError" in msg
-    assert "['reasoning']" in msg
+    class _RolloutCollectionHelper:
+        def run_examples(self, examples, head_server_config):
+            return [malformed_rollout()]
+
+    class _Tokenizer:
+        pad_token_id = 7
+        eos_token_id = 8
+
+    class _MockSelf:
+        cfg = {"backfill_failed_rollouts": True}
+        rch = _RolloutCollectionHelper()
+        head_server_config = object()
+
+    mock_self = _MockSelf()
+    modified_class = NemoGym.__ray_metadata__.modified_class
+    mock_self._postprocess_nemo_gym_to_nemo_rl_result = MethodType(
+        modified_class._postprocess_nemo_gym_to_nemo_rl_result, mock_self
+    )
+    mock_self._zero_reward_nemo_rl_result = MethodType(
+        modified_class._zero_reward_nemo_rl_result, mock_self
+    )
+
+    items = _collect_rollout_stream(
+        mock_self, [{"_rowidx": 0, "agent_ref": {"name": "agent"}}], _Tokenizer()
+    )
+
+    assert len(items) == 1
+    rowidx, result, _ = items[0]
+    assert rowidx == 0
+    assert result["full_result"]["reward"] == 0.0
+    assert result["full_result"]["nemo_rl_fallback_reason"].startswith(
+        "postprocess_failed:KeyError"
+    )
+
+
+def test_nemo_gym_malformed_response_is_masked_zero_reward():
+    class _Tokenizer:
+        pad_token_id = 7
+        eos_token_id = 8
+
+    class _MockSelf:
+        cfg = {"backfill_failed_rollouts": True}
+
+    mock_self = _MockSelf()
+    modified_class = NemoGym.__ray_metadata__.modified_class
+    mock_self._zero_reward_nemo_rl_result = MethodType(
+        modified_class._zero_reward_nemo_rl_result, mock_self
+    )
+
+    result = modified_class._postprocess_nemo_gym_to_nemo_rl_result(
+        mock_self,
+        {"response": {"output": None}, "reward": 1.0},
+        _Tokenizer(),
+    )
+
+    assert result["full_result"]["reward"] == 0.0
+    assert result["full_result"]["nemo_rl_fallback_reason"] == (
+        "malformed_response_output:NoneType"
+    )
 
 
 @pytest.mark.nemo_gym
diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py
index d402e682ee9..81f687ed8e3 100644
--- a/tests/unit/models/generation/test_vllm_generation.py
+++ b/tests/unit/models/generation/test_vllm_generation.py
@@ -1858,6 +1858,98 @@ def test_VllmAsyncGenerationWorker_replace_prefix_tokens(tokenizer):
     assert result == model_token_ids
 
 
+def test_replace_prefix_tokens_empty_model_prefix_returns_template():
+    class _T:
+        eos_token_id = 2
+
+    tokenizer = _T()
+    model_prefix_token_ids = []
+    template_prefix_token_ids = [9, 2]
+    template_token_ids = [9, 2, 33, 44]
+    result = replace_prefix_tokens(
+        tokenizer=tokenizer,
+        model_prefix_token_ids=model_prefix_token_ids,
+        template_prefix_token_ids=template_prefix_token_ids,
+        template_token_ids=template_token_ids,
+    )
+    assert result == template_token_ids
+
+
+def test_replace_prefix_tokens_missing_eos_in_template_prefix_raises():
+    class _T:
+        eos_token_id = 2
+
+        def decode(self, *args, **kwargs):
+            pass
+
+    tokenizer = _T()
+    model_prefix_token_ids = [7, 2]
+    template_prefix_token_ids = [9, 9, 9]  # no EOS inside prefix
+    template_token_ids = [9, 9, 9, 2, 10]
+    with pytest.raises(AssertionError):
+        replace_prefix_tokens(
+            tokenizer=tokenizer,
+            model_prefix_token_ids=model_prefix_token_ids,
+            template_prefix_token_ids=template_prefix_token_ids,
+            template_token_ids=template_token_ids,
+        )
+
+
+def test_replace_prefix_tokens_tokenizer_without_eos_raises():
+    class _T:
+        eos_token_id = None
+
+    tokenizer = _T()
+    with pytest.raises(AssertionError):
+        replace_prefix_tokens(
+            tokenizer=tokenizer,
+            model_prefix_token_ids=[1],
+            template_prefix_token_ids=[1, 2],
+            template_token_ids=[1, 2],
+        )
+
+
+def test_replace_prefix_tokens_uses_last_eos_in_template_prefix():
+    class _T:
+        eos_token_id = 2
+
+    tokenizer = _T()
+    model_prefix_token_ids = [100, 2]
+    template_prefix_token_ids = [9, 2, 9, 2]  # two EOS; last at idx=3
+    template_token_ids = [9, 2, 9, 2, 77, 88]
+    result = replace_prefix_tokens(
+        tokenizer=tokenizer,
+        model_prefix_token_ids=model_prefix_token_ids,
+        template_prefix_token_ids=template_prefix_token_ids,
+        template_token_ids=template_token_ids,
+    )
+    assert result == [100, 2, 77, 88]
+
+
+def test_replace_prefix_tokens_repairs_qwen_normalized_prefix():
+    """Qwen3.5 re-renders assistant history with normalized reasoning content.
+
+    The isolated prefix render differs token-wise from what the model produced;
+    EOS-count splicing must still keep the original generated tokens and resume
+    from the template after the boundary EOS.
+    """
+
+    class _T:
+        eos_token_id = 2
+
+    tokenizer = _T()
+    result = replace_prefix_tokens(
+        tokenizer=tokenizer,
+        model_prefix_token_ids=[10, 11, 2],
+        # Rendering the assistant prefix alone changed token 11 to token 99.
+        template_prefix_token_ids=[10, 99, 2],
+        # Two messages follow that assistant turn, then the generation marker.
+        template_token_ids=[10, 11, 2, 20, 2, 30, 2, 40],
+    )
+
+    assert result == [10, 11, 2, 20, 2, 30, 2, 40]
+
+
 @pytest.mark.asyncio
 async def test_vllm_http_server_correct_merged_tokens_matches_baseline(
     cluster, tokenizer