Adaptive in-flight cap from vLLM load in AsyncGRPO

[AmineDiro]

AsyncGRPO: cap in-flight rollouts on the vLLM server's KV cache usage

What this does

The rollout worker currently admits up to max_inflight_tasks requests to vLLM (default max_staleness * samples_per_step, 512 in the runs below). That number knows nothing about the server: when the KV cache cannot hold that many sequences, vLLM queues the rest and preempts running ones as they grow, so every request spends most of its life waiting inside the server before it decodes.

This PR makes the worker follow the server instead. A small control loop polls /metrics every 5 s and keeps a live cap on in-flight requests:

• Start at kv_cache_size_tokens / max_model_len, as many full-length sequences as the cache can hold, never above max_inflight_tasks.
• Cut the cap by 30 % when KV cache usage is above 0.8. The cut is latched until the pool has drained below the new cap, so one spike produces one cut.
• Grow the cap by 25 % when usage is below 0.6 and nothing is queued, once per turnover of the pool, up to max_inflight_tasks.
• Nothing is cancelled. The cap only gates admission; requests already in flight finish.

New rollout metrics: rollout/inflight_cap, rollout/kv_cache_usage, rollout/vllm_running_requests, rollout/vllm_queued_requests, rollout/vllm_preemptions_total.

It is the single-environment version of prime-rl's ConcurrencyController (PrimeIntellect-ai/prime-rl#3285). No new config: everything is read from the server.

Experiment

Sanity bed: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B on sail/Sanity-Test-R1D-1.5B, vLLM 0.29 with data parallel 2 on two H200, the trainer on two more, 128 completions per step, max_completion_length=8192, max_staleness=4, 300 steps. To make the KV cache the bottleneck, as it is for any real model size, the server was started with --kv-cache-memory-bytes 8G per replica: about 0.6 M tokens in total, or roughly 100 completions at the 6.1 k mean length this recipe produces. Everything else is identical between the two runs.

• kv8g-main: main @ 72cf37d, 512 in flight
• kv8g-cap: this branch, adaptive cap

Curves: https://huggingface.co/spaces/aminediroHF/trl-sanity

The cap tracks the server

The cap starts near 60 from the cache size, grows while usage is under 0.6 and backs off as it passes 0.8, settling between 70 and 130 requests, which is what the cache actually holds. main pushes 512 regardless: the server runs about 100 of them and queues the other 400.

What over-admission costs

metric	main	this PR
requests in flight	512, fixed	70 to 134
vLLM KV cache usage	not reported	0.41 to 0.81
vLLM preemptions	TBD from the job log	TBD
rollout latency, dispatch to last token	94 s	24 s
sample staleness, mean	3.7 policy steps	1.6
stale samples dropped	2984 of ~10.6 k consumed (28 %)	0
generated tokens/s, 2 x H200	32 k	27 k
reward, last 25 steps	0.569	0.570

Completions are the same length on both runs, so the 4x latency gap is queue time inside vLLM. At 30 s per optimizer step a sample that takes 94 s comes back three policies late, which is why main sits at the staleness ceiling and drops 28 % of what it generates. The cap run never submits more than the cache holds, every request decodes immediately, and nothing is dropped.

Raw generation throughput is about 15 % higher on main: a fuller decode batch is more efficient per token. But 28 % of those tokens are thrown away as stale, so useful throughput is higher with the cap, and the trainer waits less per step. The learning curve is the same, as it should be: the cap changes when rollouts are admitted, not what is trained on.

Why we want it

• Any model larger than this one hits the cache limit at the default in-flight count. Without the cap the symptom is a flat, late, stale pipeline that looks like a slow server.
• The right in-flight count depends on the model, the hardware, the sequence length and the prompt mix, and it changes as the policy's completions get longer. Deriving it from the server removes a knob nobody can set by hand.
• It costs nothing when the cache is not the bottleneck: on the same bed with the full 0.85 memory fraction (8 M tokens), the cap sits at the ceiling for the whole run and the curve matches main.

Not in this PR

Staleness on an oversized pool is a separate problem. With the full cache, 512 in flight still drops 42 % of samples as stale on this bed because the pool is sized by max_staleness rather than by what the trainer consumes. The cap only reacts to server load, so that needs its own change.

---

Note

Medium Risk
Changes rollout admission and weight-sync pause semantics for all AsyncGRPO runs; mis-parsed metrics fall back to the ceiling, but wrong caps could still throttle or overload vLLM.

Overview
Adds adaptive concurrency for AsyncGRPO rollouts: the worker no longer always admits up to max_inflight_tasks vLLM requests. A background _control_loop polls /metrics every few seconds and adjusts a live _inflight_cap (still bounded by max_inflight_tasks). Startup sets the cap from kv_cache_size_tokens / max_model_len when metrics are available; under load it cuts by 30% when KV usage exceeds 0.8 and grows by 25% when usage is below 0.6 with no queued requests. Admission in _generate_loop gates on this cap instead of a fixed slot pool.

Documents the behavior under a new Concurrency section and logs rollout/inflight_cap, rollout/kv_cache_usage, and related vLLM queue/running/preemption metrics.

Weight sync changes: vLLM pause now uses clear_cache=false (async GRPO and async distillation clients) so cache clears do not preempt in-flight generations. WeightTransferClient accepts compute_dtype from trainer bf16/fp16 so dtype mismatch warnings compare against compute precision, not fp32 master weights.

Includes an integration test test_inflight_cap_follows_vllm_load with a fake metrics server.

^{Reviewed by Cursor Bugbot for commit a5e14a6. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale cache across weight updates

High Severity

pause now passes clear_cache=false on every weight sync, so vLLM keeps KV and prefix-cache blocks built with the previous weights. On the merged-weight path those blocks stay reusable, and later requests, especially repeated prompts, can decode against a mix of old and new weights. The LoRA path avoids this only because each sync uses a new adapter name.

 

^{Reviewed by Cursor Bugbot for commit 3494835. Configure here.}

[AmineDiro]

I looked a bit deeper into how vLLM handles this, and I think there are a few important points.

First, in AsyncGRPO we already tag rollouts with the policy version before dispatching them, and drop samples that exceed max_staleness. So we already account for policy staleness during training. Recomputing the KV cache doesn't change the tokens or logprobs already sampled with the previous weights.

The real concern is the prefix cache for new requests. With clear_cache=false, a new request could reuse a prefix computed several policy versions ago. This isn't captured by our staleness check, so I agree there's a potential issue here.

But the performance cost of clearing the cache is huge, and I think this is the important part:

• TRL supports vLLM >=0.21. Until now, we called pause(mode="keep") without specifying clear_cache, so it defaulted to true.
• In practice, this preempts every running request at each weight sync, clears the prefix cache, and forces all requests to prefill again from scratch when resumed.
• With clear_cache=false, running requests keep their KV cache and resume decoding directly after the sync, without any extra prefill.

In our sanity run, that's around 500 requests × 3k tokens = 1.5M tokens unnecessarily re-prefilled at every sync. This explains the bursty throughput we were seeing on main.

So clear_cache=false is actually the behavior our TRL docstring described IMO.

I'd prefer keeping false. We already handle rollout staleness, and avoiding this massive amount of recomputation is important for async. @qgallouedec @kashif

[AmineDiro]

@kashif @qgallouedec

Follow-up on the review question about whether the gains come from the cap or from clear_cache=false: two more runs to isolate each change. Same setup as above, with --kv-cache-memory-bytes 8G per replica.

Run	Cap	Sync keeps cache	Ref
kv8g-main	No	No	main @ 72cf37d
kv8g-pausefix	No	Yes	532025a
kv8g-caponly	Yes	No	2e0959b
kv8g-cap	Yes	Yes	142be98

Metric	Main	Pause fix only	Cap only	Both (this PR)
Wall clock, 300 steps	138 min	129 min	123 min	117 min
vLLM preemptions, whole run	501,678	436,774	35,638	21
Requests in flight	512	512	74–163	70–162
Rollout latency	81 s	77 s	22 s	21 s
Sample staleness, mean	3.7	3.8	1.8	1.9
Stale samples dropped	34%	32%	0%	0%
Generated tokens/s	34k	36k	29k	30k
Reward, steps 251–300	0.679	0.678	0.659	0.663

The cap is what removes the queueing: 4× lower latency, half the staleness, no drops, and 15 minutes saved. Keeping the cache on its own saves 9 minutes but changes little else, because with 512 requests competing for a ~100-request cache, the server is still preempting for memory regardless.

Keeping the cache really pays off once the cap holds the server within its cache capacity. The 35,638 preemptions left in the cap-only run are roughly one per running request per sync, corresponding to prefill recomputation. With clear_cache=false, that drops to just 21.

Curves: https://huggingface.co/spaces/aminediroHF/trl-sanity (runs kv8g-*).

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 2ada061. Configure here.}