feat(sc): support router replay with TQ

docs/guides/router-replay.md

@@ -34,6 +34,12 @@ An example recipe is available at:
 examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml
 ```
 
+The native async TransferQueue path uses the SingleController entrypoint with:
+
+```text
+examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml
+```
+
 ## Validation
 
 Router Replay validation covers two end-to-end questions:
@@ -43,13 +49,14 @@ Router Replay validation covers two end-to-end questions:
 2. whether matched R3-on runs reduce train-vs-rollout mismatch relative to
    matched R3-off controls.
 
-### Trace Debugging
+### Validation and Trace Debugging
 
 Router Replay can emit JSONL traces for a small number of training steps. This
 is intended for correctness debugging, not long training runs.
 
 | Environment variable | Default | Meaning |
 | --- | --- | --- |
+| `NRL_ROUTER_REPLAY_VALIDATE` | `0` | Validate replay tensors before Megatron installs them, rejecting partially missing routes, duplicate top-k expert IDs, and out-of-range expert IDs. |
 | `NRL_R3_TRACE` | `0` | Master switch for R3 JSONL trace emission. |
 | `NRL_R3_TRACE_STEPS` | `1` | Number of training steps to trace. |
 | `NRL_R3_TRACE_SAMPLES` | `2` | Number of samples with full tensor previews. |
@@ -110,7 +117,11 @@ all returned vLLM routes are still replayed exactly.
 The fallback is intentionally route-local: it does not disable Router Replay for
 the whole batch or sample.
 
-When fallback is used, NeMo RL logs
-`r3/routed_experts_fallback_token_route_fraction`. This metric should normally
-be zero or near-zero. A nonzero value means some token routes used Megatron's
-normal router instead of replay.
+When fallback is used, the vLLM worker emits a `R3 router replay fallback:` warning
+to the run log naming the affected sample count and missing token-route count.
+Fallback should normally be absent or rare; frequent warnings mean a meaningful
+share of token routes used Megatron's normal router instead of replay.
+
+The generation backend also computes
+`r3/routed_experts_fallback_token_route_fraction`, but no training loop currently
+forwards it to the metric logger, so do not rely on it in dashboards or gates.

examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml

@@ -0,0 +1,29 @@
+# SingleController async+TQ variant of the Qwen3-30B-A3B R3 recipe.
+# Launch with:
+#   uv run examples/run_grpo_single_controller.py \
+#     --config examples/configs/recipes/llm/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml
+defaults: ./grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async.yaml
+
+# SC does not support validation and checkpointing yet.
+grpo:
+  val_period: 0
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller
+
+# TransferQueue is mandatory for the SingleController path.
+data_plane:
+  enabled: true
+
+async_rl:
+  sampler:
+    name: windowed
+    max_staleness_versions: 1
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  max_inflight_prompts: ${grpo.num_prompts_per_step}
+  max_buffered_rollouts: 64
+
+logger:
+  log_dir: logs/grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller
+  wandb:
+    name: grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -24,8 +24,10 @@
 
 from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol
 from nemo_rl.data_plane import KVBatchMeta
+from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD
 from nemo_rl.experience.interfaces import PromptGroupRecord
 from nemo_rl.experience.payload import pack_payload, record_to_train_batch
+from nemo_rl.utils.r3_trace import trace_rollout_payload
 
 
 # Classes with @ray.remote can't be inherited from, so we split the implementation out.
@@ -648,10 +650,12 @@ def __init__(
         partition_id: str,
         *,
         pad_value_dict: Mapping[str, int],
+        require_routed_experts: bool = False,
     ):
         self._dp_client = dp_client
         self._partition_id = partition_id
         self._pad_value_dict = dict(pad_value_dict)
+        self._require_routed_experts = require_routed_experts
         self.meta_list: list[Optional[KVBatchMeta]] = []
         self.start_weight_list: list[int] = []
         self.end_weight_list: list[int] = []
@@ -708,6 +712,7 @@ async def commit(
 
         Raises:
             ValueError: group_id has no live slot (removed or never reserved).
+            RuntimeError: router replay is enabled but the payload has no routes.
         """
         # Precondition: reserve() must have registered this group_id. Raise
         # before any side effects so a stray commit doesn't leak orphan DP rows.
@@ -720,6 +725,14 @@ async def commit(
         sample_ids, fields, tags = pack_payload(
             train_batch, weight_version=start_weight_version, group_id=group_id
         )
+        if self._require_routed_experts and ROUTED_EXPERTS_FIELD not in fields:
+            raise RuntimeError(
+                "policy.router_replay.enabled=true requires routed_experts in "
+                "the SingleController rollout payload, but payload packing did "
+                "not produce that field. Check vLLM routed-expert capture and "
+                "the async message-log flattening path."
+            )
+        trace_rollout_payload(keys=sample_ids, data=train_batch)
         try:
             await self._call_dp(
                 "put_samples",

nemo_rl/algorithms/grpo.py

@@ -86,6 +86,7 @@
 )
 from nemo_rl.experience.rollouts import (
     EffortLevelsConfig,
+    backfill_missing_routed_experts,
     get_nemo_gym_thinking_tags,
     run_async_multi_turn_rollout,
     run_multi_turn_rollout,
@@ -1740,12 +1741,16 @@ def add_grpo_token_loss_masks_and_generation_logprobs(
     generated assistant messages have generation_logprobs, so use that field as the
     trainable-token marker. This function mutates each message in-place by adding a
     token_loss_mask and, when missing, a zero-valued generation_logprobs tensor.
+    Router-replay routes get the same treatment via
+    :func:`backfill_missing_routed_experts`, so every per-token field is defined
+    for every tokenized message before the batch is flattened.
 
     Args:
         message_logs: Batch of tokenized message logs. Each message must contain a
             ``role`` and ``token_ids`` field. Messages that already contain
             ``generation_logprobs`` are treated as rollout-generated messages.
     """
+    backfill_missing_routed_experts(message_logs)
     for message_log in message_logs:
         for message in message_log:
             role = cast(str, message["role"])
@@ -2948,6 +2953,10 @@ def grpo_train(
                     # Save baseline for logging (before deletion)
                     baseline_for_log = baseline.clone()
 
+                    # Must precede prompt extraction: it reuses the same message
+                    # dicts, so this also protects the prompt flatten below.
+                    backfill_missing_routed_experts(repeated_batch["message_log"])
+
                     # Extract original prompt messages using the length field
                     # This correctly handles multi-turn prompts that contain assistant messages
                     initial_prompt_message_logs = extract_initial_prompt_messages(
@@ -3860,9 +3869,12 @@ def async_grpo_train(
         master_config.data_plane or {}
     ).get("enabled", False):
         raise NotImplementedError(
-            "policy.router_replay.enabled=true with async GRPO is currently "
-            "supported only when data_plane.enabled=false. Async + TQ support "
-            "has not been merged yet."
+            "policy.router_replay.enabled=true with async GRPO on this "
+            "entrypoint is supported only when data_plane.enabled=false. For "
+            "async + TransferQueue, use the SingleController entrypoint: "
+            "examples/run_grpo_single_controller.py with e.g. "
+            "examples/configs/recipes/llm/"
+            "grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml"
         )
 
     if master_config.grpo["async_grpo"]["max_trajectory_age_steps"] > 1:
@@ -4353,6 +4365,10 @@ def async_grpo_train(
 
                 print("▶ Processing rewards...")
                 with timer.time("reward_calculation"):
+                    # Must precede prompt extraction: it reuses the same message
+                    # dicts, so this also protects the prompt flatten below.
+                    backfill_missing_routed_experts(repeated_batch["message_log"])
+
                     # Extract original prompt messages using the length field
                     # This correctly handles multi-turn prompts that contain assistant messages
                     initial_prompt_message_logs = extract_initial_prompt_messages(

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -432,6 +432,7 @@ def setup_single_controller(
         dp_client,
         partition_id=partition_id,
         pad_value_dict={"token_ids": pad_id, "input_ids": pad_id},
+        require_routed_experts=router_replay_enabled(policy_config),
     )
     rollout_manager = RolloutManager(
         tokenizer=tokenizer,

nemo_rl/experience/payload.py

@@ -23,6 +23,7 @@
 
 from nemo_rl.data_plane.codec import pack_jagged_fields
 from nemo_rl.data_plane.column_io import TOKEN_ALIGNED_FIELDS
+from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.experience.interfaces import PromptGroupRecord
 
@@ -40,7 +41,7 @@ def record_to_train_batch(
 
     Returns:
         BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask,
-        sample_mask, prompt_ids_for_adv, and total_reward.
+        sample_mask, prompt_ids_for_adv, total_reward, and optional routed_experts.
     """
     # Lazy imports: grpo and llm_message_utils transitively pull
     # experience.rollouts, so importing at module top risks a cycle.
@@ -49,6 +50,7 @@ def record_to_train_batch(
         extract_initial_prompt_messages,
     )
     from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+    from nemo_rl.experience.rollouts import backfill_missing_routed_experts
 
     completions = record.completions
     n = len(completions)
@@ -58,6 +60,11 @@ def record_to_train_batch(
     prompt_token_count = sum(len(m["token_ids"]) for m in record.prompt)
     prompt_lengths = torch.full((n,), prompt_token_count, dtype=torch.long)
 
+    # Must precede the prompt extraction: it reuses the same message dicts, so
+    # backfilling here also covers the prompt flatten below. Doing it only inside
+    # add_grpo_token_loss_masks_and_generation_logprobs would be too late.
+    backfill_missing_routed_experts(message_logs)
+
     prompt_message_logs = extract_initial_prompt_messages(message_logs, prompt_lengths)
     prompt_flat, _ = batched_message_log_to_flat_message(
         prompt_message_logs,
@@ -75,17 +82,18 @@ def record_to_train_batch(
     )
     sample_mask = torch.ones(n, dtype=torch.float32)
 
-    return BatchedDataDict[Any](
-        {
-            "input_ids": flat["token_ids"],
-            "input_lengths": input_lengths,
-            "generation_logprobs": flat["generation_logprobs"],
-            "token_mask": flat["token_loss_mask"],
-            "sample_mask": sample_mask,
-            "prompt_ids_for_adv": prompt_flat["token_ids"],
-            "total_reward": total_reward,
-        }
-    )
+    train_data: dict[str, Any] = {
+        "input_ids": flat["token_ids"],
+        "input_lengths": input_lengths,
+        "generation_logprobs": flat["generation_logprobs"],
+        "token_mask": flat["token_loss_mask"],
+        "sample_mask": sample_mask,
+        "prompt_ids_for_adv": prompt_flat["token_ids"],
+        "total_reward": total_reward,
+    }
+    if ROUTED_EXPERTS_FIELD in flat:
+        train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD]
+    return BatchedDataDict[Any](train_data)
 
 
 def pack_payload(

nemo_rl/experience/rollout_manager.py

@@ -27,7 +27,13 @@
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.experience.interfaces import Completion, PromptGroupRecord
 from nemo_rl.experience.metric_utils import calculate_single_metric, pct
-from nemo_rl.experience.rollouts import _tensorize_by_key, calculate_rewards
+from nemo_rl.experience.rollouts import (
+    _attach_routed_experts_to_message_log_prefix,
+    _dummy_routed_experts_for_tokens,
+    _find_routed_experts_template,
+    _tensorize_by_key,
+    calculate_rewards,
+)
 from nemo_rl.models.generation.interfaces import (
     GenerationConfig,
     GenerationDatumSpec,
@@ -216,13 +222,17 @@ async def _run_single_rollout(
                     tokenized_obs = torch.empty(0, dtype=tokenized_obs.dtype)
                 truncated = True
 
-            current_message_log.append(
-                {
-                    "role": env_output.observations[0]["role"],
-                    "content": env_obs_content,
-                    "token_ids": tokenized_obs,
-                }
-            )
+            env_message: dict[str, Any] = {
+                "role": env_output.observations[0]["role"],
+                "content": env_obs_content,
+                "token_ids": tokenized_obs,
+            }
+            routed_template = _find_routed_experts_template(current_message_log)
+            if routed_template is not None:
+                env_message["routed_experts"] = _dummy_routed_experts_for_tokens(
+                    tokenized_obs, routed_template
+                )
+            current_message_log.append(env_message)
 
             # Update token counts
             env_token_count += len(tokenized_obs)
@@ -303,6 +313,17 @@ async def _generate_response(
             assistant_message["generation_logprobs"] = output["logprobs"][
                 0, input_len:total_len
             ]
+        if "routed_experts" in output:
+            routed_experts = output["routed_experts"][0]
+            prefix_length = _attach_routed_experts_to_message_log_prefix(
+                message_log, routed_experts
+            )
+            if prefix_length != input_len:
+                raise RuntimeError(
+                    "message_log token length does not match generation input_length "
+                    f"({prefix_length} != {input_len})."
+                )
+            assistant_message["routed_experts"] = routed_experts[input_len:total_len]
 
         # Calculate generation metrics
         gen_metrics: dict[str, Any] = {}

nemo_rl/experience/rollouts.py

@@ -50,6 +50,7 @@
 from nemo_rl.experience.interfaces import NEMO_GYM_TASK_INDEX_KEY
 from nemo_rl.experience.metric_utils import calculate_single_metric, pct
 from nemo_rl.models.generation.interfaces import (
+    ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL,
     GenerationConfig,
     GenerationDatumSpec,
     GenerationInterface,
@@ -144,6 +145,50 @@ def _dummy_routed_experts_for_tokens(
     )
 
 
+def backfill_missing_routed_experts(
+    message_logs: Sequence[list[dict]],
+) -> None:
+    """Give every tokenized message a ``routed_experts`` row, in place.
+
+    Routes are attached only where generation ran, so a trajectory whose first
+    turn raised (or a turn whose routes vLLM could not return) leaves messages
+    without the field while its siblings have it. Flattening then either stacks
+    ragged ranks or silently concatenates a short column, so fill the gaps with
+    the all--1 missing-route sentinel: Megatron routes those tokens with its own
+    router, which is the honest answer for tokens no capture covered.
+
+    No-op when the batch carries no routes at all — that is the router-replay-off
+    case, and on the TQ paths the producer-side guard must still see the field
+    missing so it can report a capture failure.
+    """
+    template = None
+    for message_log in message_logs:
+        template = _find_routed_experts_template(message_log)
+        if template is not None:
+            break
+    if template is None:
+        return
+    if template.dim() != 3:
+        raise ValueError(
+            "routed_experts messages must have shape [tokens, layers, topk], "
+            f"got {tuple(template.shape)}"
+        )
+
+    for message_log in message_logs:
+        for msg in message_log:
+            token_ids = msg.get("token_ids")
+            if not isinstance(token_ids, torch.Tensor):
+                continue
+            if isinstance(msg.get("routed_experts"), torch.Tensor):
+                continue
+            msg["routed_experts"] = torch.full(
+                (int(token_ids.shape[0]), template.shape[1], template.shape[2]),
+                ROUTED_EXPERTS_MISSING_ROUTE_SENTINEL,
+                dtype=template.dtype,
+                device=template.device,
+            )
+
+
 class EffortLevelsConfig(BaseModel, extra="allow"):
     """Controls length-based reward shaping for low-effort prompts.
 

nemo_rl/experience/sync_rollout_actor.py

@@ -80,8 +80,12 @@ def _flatten_rollout_message_log_for_tq(
         extract_initial_prompt_messages,
     )
     from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message
+    from nemo_rl.experience.rollouts import backfill_missing_routed_experts
 
     pad = {"pad_value_dict": {"token_ids": pad_token_id}}
+    # Must precede the prompt extraction: it reuses the same message dicts, so
+    # backfilling here also covers the prompt flatten below.
+    backfill_missing_routed_experts(message_logs)
     prompt_message_logs = extract_initial_prompt_messages(
         message_logs,
         prompt_lengths,