fix(async): do not let the vLLM prefix cache go stale across weight updates

[linj-glitch]

What does this PR do ?

Stops the vLLM prefix cache from silently going stale across weight updates in async GRPO/PPO.

With async rollouts and vLLM prefix caching enabled (the default on Ampere+), recompute_kv_cache_after_weight_updates=false did more than keep in-flight requests on their pre-update KV: it also left the prefix cache untouched across refits. Any later request that shares a prefix with an earlier one (same prompt, shared system prompt, a multi-sample group spanning a refit) was served KV/state computed by whichever weights first filled that block, with no bound on how old those weights were. The trainer scores with the current weights, so the train-vs-rollout logprob mismatch grew step over step within a job and reset only on engine restart. The docs described the flag as affecting in-flight requests only.

Evidence

Nemotron-H MoE (Nano 3.5), async GRPO, in_flight_weight_updates=true, flag false, prefix caching default-on, ~49k-token prompts where same-instance rows share ~85% of their context and recur 3–20 steps apart. Logs show Successfully reset prefix cache exactly at engine start and at generation finish, never across the 17 refits in between, so every later request sharing a prefix with an earlier one was served KV/Mamba state computed by weights up to 17 steps old while the trainer scored it with current weights. A controlled A/B from the same checkpoint with enable_prefix_caching=false confirmed the mechanism is not what drove the growth of that run's masked-sequence count (that turned out to be a policy-side format collapse, unrelated to caching), so this PR is a correctness fix for silent staleness rather than the cause of a specific regression. The staleness itself is unbounded by construction and the docs described the flag as affecting in-flight requests only.

Changes

• AsyncTrajectoryCollector.resume_after_refit: refits that drained generation first (no in-flight pause requested) now always invalidate the KV/prefix cache after the update, regardless of the flag. Nothing is using the cache at that point, so this is free. The flag now only decides what happens to requests that were in flight during an in-flight refit (preempt and recompute vs. keep pre-update KV, Magistral-style).
• AsyncTrajectoryCollector.prepare_for_refit: in-flight refit with the flag false and prefix caching enabled prints a one-time warning explaining the unbounded staleness and how to fix it (recompute_kv_cache_after_weight_updates=true or vllm_cfg.enable_prefix_caching=false).
• Default of recompute_kv_cache_after_weight_updates flipped to true in AsyncGRPOConfig and the PPO async config. Recipes that set it explicitly are unchanged. Cost for in-flight users who relied on the old default: in-flight requests re-prefill once per refit (vLLM pause_generation(mode="keep", clear_cache=True) preempts them and resets the block pool).
• docs/guides/async-grpo.md: describe both effects of the flag; example snippets use true.
• Tests (tests/unit/algorithms/test_async_utils.py): drained refit invalidates even with the flag false; in-flight + false warns exactly once and keeps the cache; no warning when prefix caching is off. Existing keep-pause and legacy-fallback tests unchanged.

If maintainers prefer not to change the default for the performance recipes, the first two changes plus the docs stand alone; the default flip is the last commit hunk in grpo.py/ppo.py.

Megatron generation backend

The second commit keeps the existing kv_cache_management_mode consistency check intact under the new default: when the flag is left unset and policy.generation.backend=megatron, it is derived from the engine-side mode (recompute → true, persist/offload → false) instead of erroring. Explicit disagreement still raises as before.

Issues

Root-caused while investigating growing train/inference mismatch in a long-context async GRPO run; no existing issue.

Usage

grpo:
  async_grpo:
    in_flight_weight_updates: true
    recompute_kv_cache_after_weight_updates: true   # new default

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (3 added, 1 rewritten in tests/unit/algorithms/test_async_utils.py)
• Did you run the unit tests and functional tests locally? tests/unit/algorithms/test_async_utils.py: 118 passed; tests/unit/algorithms/test_grpo.py + test_ppo.py -k 'async or recompute or kv_cache or refit or config or setup': 107 passed; tests/unit/single_controller/test_refit_recovery.py + tests/unit/models/generation/test_swe1_dynamo_config.py: 37 passed (run inside the NeMo RL container).
• Did you add or update any necessary documentation? (docs/guides/async-grpo.md)

Additional Information

• Found while root-causing a growing train-vs-rollout logprob mismatch; the cache staleness was one of the candidates and is real, though the run's growth was ultimately a policy-side effect.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.