feat: pause SingleController refits and enable KV recomputation by default

docs/guides/async-grpo.md

@@ -43,7 +43,7 @@ grpo:
     max_trajectory_age_steps: 1  # Maximum age, in training steps, for trajectories
     max_generation_failures: 0  # Consecutive worker failures to tolerate
     in_flight_weight_updates: false  # Enable for faster weight synchronization
-    recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates
+    recompute_kv_cache_after_weight_updates: true # Recompute KV after weight updates (default)
 ```
 
 ### Complete Example Config
@@ -72,7 +72,7 @@ grpo:
     max_trajectory_age_steps: 1
     max_generation_failures: 0  # Consecutive worker failures to tolerate
     in_flight_weight_updates: false  # Enable for faster weight synchronization
-    recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates
+    recompute_kv_cache_after_weight_updates: true # Recompute KV after weight updates (default)
 
 cluster:
   num_nodes: 2
@@ -221,9 +221,9 @@ If no `replay_buffer.pt` file is found in the latest checkpoint directory, train
 
 3. **Resource Allocation**: Ensure sufficient GPU memory for both the training and generation clusters
 
-4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain; SGLang is in this group today and instead relies on the pause its own weight synchronizer performs around the transfer. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine.
+4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. SGLang acknowledges these hooks through its weight synchronizer, which pauses, flushes KV, transfers weights, and resumes generation; its configured pause mode determines whether requests are requeued or aborted. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine.
 
-5. **Recompute KV Cache After Weight Updates**: Set `recompute_kv_cache_after_weight_updates: true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. Other refit paths keep their existing post-update invalidation behavior. When false, in-flight requests retain their pre-update KV cache. On the Megatron generation backend, this must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`; setup errors on a mismatch.
+5. **Recompute KV Cache After Weight Updates**: `recompute_kv_cache_after_weight_updates` defaults to `true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. This repeats prefill for those requests and can reduce throughput. SGLang always flushes KV inside its weight synchronizer before the transfer, regardless of this flag; on its in-flight path, the collector skips a second flush after generation resumes. Other refit paths keep their existing post-update invalidation behavior. Set the flag to `false` to retain pre-update KV caches and avoid this recomputation cost on backends that support doing so. On the Megatron generation backend, the flag must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`: use `true` with `recompute`, or explicitly set `false` with `persist`/`offload`; setup errors on a mismatch.
 
 ## Why Importance Sampling Correction Is Required for Async
 

docs/guides/ppo.md

@@ -276,7 +276,7 @@ ppo:
     max_trajectory_age_steps: 1
     warmup_generation_lead_steps: null
     in_flight_weight_updates: false
-    recompute_kv_cache_after_weight_updates: false
+    recompute_kv_cache_after_weight_updates: true  # Default; false avoids repeated prefill but retains old KV.
     drop_incomplete_targets_on_restore: false
 
   adv_estimator:

docs/guides/single-controller.md

@@ -422,7 +422,9 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th
 2. **Actor startup**: `SingleControllerActor` launches `_rollout_pump` and `_train_pump` concurrently as asyncio tasks; both share the same `TQReplayBuffer` and `StalenessSampler`.
 3. **Rollout pump loop**: `sampler.admit` gates dispatch against the current trainer version (returning a `target_step` for `in_order`); the pump then reserves a buffer slot, drives `RolloutManager.generate_and_push`, and commits with the observed `start_weight` / `end_weight`.
 4. **Train pump loop**: `sampler.evict` removes out-of-window groups, and `sampler.select` chooses the next chunk. For PPO, `_value_stage` computes critic values and `_advantage_stage` computes advantages. TQPolicy then updates the policy: full-batch PPO uses `ppo.ppo_epochs` updates, while streaming PPO accumulates gradients across chunks for one update. Streaming PPO refits the updated policy before `_value_train_epochs` trains the critic on the assembled full batch for `ppo.critic_ppo_epochs`. GRPO uses one TQPolicy optimizer step per RL step.
-5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version. In streaming PPO, it does the weight sync as soon as the policy training is finished, which is before the critic model training starts. This allows new batch generation with the updated policy weight overlapping with the critic training epochs.
+5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior. In streaming PPO, it does the weight sync as soon as the policy training is finished, which is before the critic model training starts. This allows new batch generation with the updated policy weight overlapping with the critic training epochs.
+
+   For non-colocated vLLM, freezing a request does not stop its end-to-end generation timeout. Size `NRL_VLLM_ASYNC_TIMEOUT_SECONDS` and, when the generation router is enabled, `async_rl.generation_router.backend_timeout_s` to cover the longest expected generation plus time spent paused for refit, including membership recovery and a possible retry. With a finite `async_rl.generation_fleet_health.refit_timeout_s`, pause, each transfer attempt, and resume each use that timeout plus a 60 s grace; these are separate call budgets, not a single total refit budget. When `refit_timeout_s` is `null`, pause and resume retain a 360 s budget each, but the collective transfer wait is unbounded, so no finite worst-case total pause duration is guaranteed.
 
 ## Relation to Legacy Async GRPO
 
@@ -449,7 +451,7 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous-
 | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async |
 | `max_trajectory_age_steps: N` | `sampler.name: in_order` with `sampler.max_lookahead_versions: N` |
 | `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress |
-| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same) |
+| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same, defaults to `true` in both loops); when vLLM pauses natively, preserved requests recompute their KV cache after refit, repeating prefill. Set `false` to retain pre-update KV. For Megatron generation, pair `true` with `kv_cache_management_mode: recompute`, or `false` with `persist`/`offload`. |
 | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) |
 | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run |
 | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` |

examples/configs/grpo_math_1B.yaml

@@ -64,7 +64,7 @@ grpo:
     # when transient generation errors are expected and acceptable to drop.
     max_generation_failures: 0
     in_flight_weight_updates: false # Set to true to enable in-flight weight updates
-    recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates.
+    recompute_kv_cache_after_weight_updates: true # Set false to retain pre-update KV caches (avoids repeated prefill).
 
 # Reward-zeroing penalties applied to NeMo-Gym rollout results.
 reward_penalties:
@@ -506,7 +506,7 @@ policy:
       prefix_cache_ttl_seconds: 300.0  # MCore default coordinator cache-entry lifetime.
       prefix_caching_routing_alpha: 1.0  # MCore default load penalty for prefix-affinity routing.
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
-      kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute".
+      kv_cache_management_mode: "recompute"  # Must agree with recompute_kv_cache_after_weight_updates. Options: "persist", "offload", "recompute".
       async_sched_mode: "async"  # Dynamic-batching scheduler mode. Options: "legacy", "async". "async" hides scheduler CPU time behind GPU compute.
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -27,7 +27,7 @@ async_rl:
     # How far generation may run ahead of the trainer.
     max_lookahead_versions: 1
   # Recompute generation KV caches after each weight update.
-  recompute_kv_cache_after_weight_updates: false
+  recompute_kv_cache_after_weight_updates: true
   # Min ready groups the streaming trainer waits for before dispatching a batch.
   min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
   # Cap on in-flight generate_and_push calls in the rollout pump.

examples/configs/ppo_math_1B.yaml

@@ -35,7 +35,7 @@ ppo:
     # Requires policy_training_start_step > 0. null uses max_trajectory_age_steps.
     warmup_generation_lead_steps: null
     in_flight_weight_updates: false
-    recompute_kv_cache_after_weight_updates: false
+    recompute_kv_cache_after_weight_updates: true
     # true discards partial restored rows and fills from subsequent prompts.
     drop_incomplete_targets_on_restore: false
 

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -39,7 +39,7 @@ async_rl:
     # throughout. Buffer capacity below must cover the wider of the two.
     warmup_lookahead_versions: null
   # Recompute generation KV caches after each weight update.
-  recompute_kv_cache_after_weight_updates: false
+  recompute_kv_cache_after_weight_updates: true
   # Full-batch default. Reduce this to stream forwards, GAE, and policy backward;
   # requires ppo_epochs=1. normalize_advantages=true then normalizes per chunk.
   # Value training stays full batch.

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml

@@ -23,6 +23,7 @@ policy:
     backend: megatron
     refit_transport: mcore
     mcore_generation_config:
+      kv_cache_management_mode: recompute
       refit_backend: nccl
       tensor_model_parallel_size: 4
       expert_model_parallel_size: 4

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml

@@ -19,6 +19,7 @@ policy:
     backend: megatron
     refit_transport: mcore
     mcore_generation_config:
+      kv_cache_management_mode: recompute
       refit_backend: nccl
       expert_model_parallel_size: 4
     colocated:

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-async-lag4.yaml

@@ -18,4 +18,3 @@ policy:
     refit_transport: mcore
     mcore_generation_config:
       refit_backend: nccl
-      kv_cache_management_mode: recompute

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-async-lag4.yaml

@@ -18,4 +18,3 @@ policy:
     refit_transport: mcore
     mcore_generation_config:
       refit_backend: nccl
-      kv_cache_management_mode: recompute

examples/configs/recipes/llm/ppo-deepseek-r1-distill-qwen-7b-8n4g-megatron-noncolocated-async-single-controller.yaml

@@ -89,7 +89,6 @@ cluster:
 async_rl:
   sampler:
     warmup_lookahead_versions: 2
-  recompute_kv_cache_after_weight_updates: true
   max_inflight_prompts: ${mul:${ppo.num_prompts_per_step}, 3}
   max_buffered_rollouts: ${mul:${ppo.num_prompts_per_step}, 3}
 checkpointing:

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml

@@ -9,6 +9,7 @@ grpo:
   async_grpo:
     enabled: true
     in_flight_weight_updates: true
+    recompute_kv_cache_after_weight_updates: false
 policy:
   model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
   is_vlm: true

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml

@@ -12,7 +12,7 @@ async_rl:
   sampler:
     name: in_order
     max_lookahead_versions: 1
-  recompute_kv_cache_after_weight_updates: false
+  recompute_kv_cache_after_weight_updates: true
   min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
   max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2}
   max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2}

examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml

@@ -43,6 +43,7 @@ policy:
     backend: "megatron"
     refit_transport: "mcore"
     mcore_generation_config:
+      kv_cache_management_mode: recompute
       refit_backend: nccl
       # NeMo-Gym drives rollouts through the engine's OpenAI server.
       expose_http_server: true

nemo_rl/algorithms/grpo.py

@@ -297,7 +297,7 @@ class AsyncGRPOConfig(BaseModel, extra="allow"):
     # without waiting for the pending generations to finish.
     in_flight_weight_updates: bool = False
     # Recomputes the KV cache after weight updates.
-    recompute_kv_cache_after_weight_updates: bool = False
+    recompute_kv_cache_after_weight_updates: bool = True
 
 
 class RewardPenaltyTokenIdsConfig(BaseModel, extra="allow"):
@@ -1352,13 +1352,15 @@ def _spinup_nemo_gym(base_urls, model_name, policy_api_key=None):
         policy_config["megatron_cfg"]["train_iters"] = total_train_iters
 
     # Megatron generation expresses recompute-after-refit engine-side via
-    # `kv_cache_management_mode="recompute"`; the loop-level flag must agree.
-    if generation_config["backend"] == "megatron":
-        async_grpo_config = grpo_config.async_grpo
-        recompute_kv_cache = bool(
-            async_grpo_config is not None
-            and async_grpo_config.recompute_kv_cache_after_weight_updates
-        )
+    # `kv_cache_management_mode="recompute"`; the async loop's flag must agree.
+    # Synchronous GRPO does not consume this flag.
+    async_grpo_config = grpo_config.async_grpo
+    if (
+        generation_config["backend"] == "megatron"
+        and async_grpo_config is not None
+        and async_grpo_config.enabled
+    ):
+        recompute_kv_cache = async_grpo_config.recompute_kv_cache_after_weight_updates
         kv_cache_mode = generation_config["mcore_generation_config"][
             "kv_cache_management_mode"
         ]

nemo_rl/algorithms/ppo.py

@@ -158,7 +158,7 @@ class AsyncPPOConfig(BaseModel, extra="allow"):
     # Allows weight updates while rollout requests are still in flight.
     in_flight_weight_updates: bool = False
     # Recomputes the KV cache after weight updates.
-    recompute_kv_cache_after_weight_updates: bool = False
+    recompute_kv_cache_after_weight_updates: bool = True
     # Drops partial restored targets; replacement rollouts use subsequent prompts.
     drop_incomplete_targets_on_restore: bool = False