Add reset_running_requests to reset_prefix_cache call for async RL - #1824
Conversation
Signed-off-by: SumanthRH <sumanthrh@anyscale.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an optional "reset_running_requests" parameter (defaulting to False) to the "reset_prefix_cache" method across various inference engines, clients, and servers, allowing workers to optionally reset in-flight requests during weight synchronization. Feedback on these changes highlights potential AttributeError risks in "vllm_server.py" and "vllm_server_actor.py" if the parsed JSON payload is not a dictionary, and notes that the synchronous "reset_prefix_cache" implementation in "BaseVLLMInferenceEngine" is redundant and violates the asynchronous interface contract.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| clear_kv_cache_on_weight_sync: bool = False | ||
| """Whether or not to clear the KV cache on weight sync. Defaults to False. | ||
| If False, we reuse KV cache from stale policies during generation | ||
| (avoids recomputation at the cost of using slightly stale KV cache). | ||
| """ |
There was a problem hiding this comment.
False by default. We will run more ablations on throughput vs accuracy for clear_kv_cache_on_weight_sync=True soon.
| """Whether or not to clear the KV cache on weight sync. Defaults to True, matching synchronous RL. | ||
| Set to False for fully async training to reuse KV cache from stale policies during generation | ||
| (avoids recomputation at the cost of using slightly stale KV cache).""" | ||
| clear_kv_cache_on_weight_sync: bool = False |
There was a problem hiding this comment.
ok unfortunately, currently, this default cannot be false, since we currently don't have a conditional in the workers on whether or not we are doing fully async training or just sync
so changing this default to false would make it so that kv cache isn't cleared for sync training as well. We should improve how this is gated but per charlie's comment on my pr: #1798 (comment) we didn't want to gate on max_staleness > 0 in case the async trainer is used with no staleness for testing.
There was a problem hiding this comment.
Fixed now! Thanks !
I am explictly avoiding the alternative of clear_cache=True for pause_generation because that is a deprecated param in vLLM - we will remove it soon
Signed-off-by: SumanthRH <sumanthrh99@gmail.com>
Signed-off-by: SumanthRH <sumanthrh99@gmail.com>
2 participants
What does this PR do?
Adds
reset_running_requeststoreset_prefix_cachefor async RL.This unifies the behaviour of
reset_prefix_cachefor sync and async RLCurrently, we call
reset_prefix_cachewithout any arguments here. In vLLM, this is finally passed to the scheduler in the engine core process of each GPU worker:https://github.com/vllm-project/vllm/blob/635c38338afe132f9555ebf4a7c8ac7dfb05b2a0/vllm/v1/core/sched/scheduler.py#L2154-L2190
if there are running requests, nothing is done, as can be seen in the KV Cache Block pool implementation here. Thus, even if
clear_kv_cache_on_weight_syncis set , stale KV cache is not cleared after a weight update with async RL right now, except at epoch boundaries.This PR adds
reset_running_requeststo clear kv cache for running requests.I did a quick test on the
fully_async_gsm8k.shscript.With clear cache (purple), TTFT shoots up when there are running requests as expected.