fix(inference): reset requests when discarding KV cache

.claude/docs/weight_sync.md

@@ -222,7 +222,7 @@ The `sharded_rdt_*.py` files are vendored from the vLLM PR: github.com/vllm-proj
 Non-colocated normally keeps the engine fully awake and does `pause_generation → broadcast → resume_generation`. The opt-in `generator.inference_engine.offload_kv_for_weight_sync` flag sleeps the engine (freeing the KV cache from GPU) *during* the sync so `gpu_memory_utilization` can be pushed higher (no need to keep KV cache resident alongside the weight-transfer scratch buffers). It turns on `enable_sleep_mode` (via `inference_servers/utils.py`). Requires non-colocated and non-LoRA. Orchestrated in `WorkerDispatch.save_weights_for_sampler`; the flow depends on the trainer:
 
 - **Synchronous trainer** (`fully_async.enabled=false`): generation is complete at sync time, so there are no in-flight requests. A plain `sleep() → wake_up(["weights"]) → broadcast → wake_up(["kv_cache"])` (the same three-phase pattern colocated uses) is enough — the standard `/sleep`+`/wake_up` endpoints discard the KV cache and free the memory.
-- **Fully-async trainer** (`fully_async.enabled=true`): generation overlaps the sync, so `pause_generation` (KEEP) freezes in-flight requests, then the allocator is driven directly (see below) so the scheduler is **not** resumed on the weights wake. The KV cache is offloaded to CPU and restored so frozen requests resume with no abort or prefill recompute — **unless** `clear_kv_cache_on_weight_sync=true`, in which case the broadcast resets the prefix cache anyway, so the KV is discarded (skipping the CPU copy) rather than offloaded.
+- **Fully-async trainer** (`fully_async.enabled=true`): generation overlaps the sync, so `pause_generation` (KEEP) freezes in-flight requests, then the allocator is driven directly (see below) so the scheduler is **not** resumed on the weights wake. The KV cache is offloaded to CPU and restored so frozen requests resume with no abort or prefill recompute — **unless** `clear_kv_cache_on_weight_sync=true`, in which case the broadcast resets running requests so their KV is recomputed after wake. This reset also runs when prefix caching is disabled: running requests still hold KV. The KV is discarded (skipping the CPU copy) rather than offloaded.
 
 The fully-async path is driven entirely from SkyRL — no vLLM patch. It deliberately avoids the `/sleep`+`/wake_up` HTTP endpoints (which route through `EngineCore.sleep`, force-clearing the prefix cache and preempting every running request at level ≥ 1). Instead it drives the per-worker `CuMemAllocator` directly via two `NewInferenceWorkerWrap` methods invoked over `/collective_rpc`:
 

skyrl/backends/skyrl_train/weight_sync/cache.py

@@ -0,0 +1,16 @@
+"""Cache invalidation policy shared by training backends during weight sync."""
+
+
+def should_reset_kv_cache(
+    *,
+    enable_prefix_caching: bool,
+    fully_async: bool,
+    clear_kv_cache_on_weight_sync: bool,
+) -> bool:
+    """Whether weight sync must invalidate cached KV, including running requests.
+
+    Synchronous training only needs to invalidate reusable prefix blocks.
+    Async training can keep requests in flight, whose KV exists even when
+    prefix caching is disabled.
+    """
+    return clear_kv_cache_on_weight_sync if fully_async else enable_prefix_caching

skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py

@@ -30,6 +30,7 @@
     WeightChunk,
     WeightExtractor,
 )
+from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache
 from skyrl.backends.skyrl_train.weight_sync.weight_extractor_utils import (
     yield_module_grouped_chunks,
 )
@@ -295,13 +296,13 @@ async def broadcast_to_inference_engines(
     ):
         if inference_engine_client is None:
             inference_engine_client = self._weight_sync_inference_client
-        use_prefix_cache = inference_engine_cfg.enable_prefix_caching
         generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype)
         cache_reset_task = None
         sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset
-        # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set
-        reset_prefix_cache: bool = use_prefix_cache and (
-            not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync
+        reset_prefix_cache = should_reset_kv_cache(
+            enable_prefix_caching=inference_engine_cfg.enable_prefix_caching,
+            fully_async=self.cfg.fully_async.enabled,
+            clear_kv_cache_on_weight_sync=self.cfg.fully_async.clear_kv_cache_on_weight_sync,
         )
         send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache}
 

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -58,6 +58,7 @@
     WeightChunk,
     WeightExtractor,
 )
+from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache
 from skyrl.backends.skyrl_train.workers.megatron.adapter_store import (
     AdapterStore,
     LoraSignature,
@@ -1525,13 +1526,13 @@ async def broadcast_to_inference_engines(
     ):
         if inference_engine_client is None:
             inference_engine_client = self._weight_sync_inference_client
-        use_prefix_cache = inference_engine_cfg.enable_prefix_caching
         generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype)
         cache_reset_task = None
         sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset
-        # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set
-        reset_prefix_cache: bool = use_prefix_cache and (
-            not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync
+        reset_prefix_cache = should_reset_kv_cache(
+            enable_prefix_caching=inference_engine_cfg.enable_prefix_caching,
+            fully_async=self.cfg.fully_async.enabled,
+            clear_kv_cache_on_weight_sync=self.cfg.fully_async.clear_kv_cache_on_weight_sync,
         )
         send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache}
 

skyrl/train/config/config.py

@@ -967,6 +967,8 @@ class FullyAsyncConfig(BaseConfig):
     the epoch ends."""
     clear_kv_cache_on_weight_sync: bool = False
     """Whether or not to clear the KV cache on weight sync. Defaults to False.
+    If True, running requests recompute their KV after the sync, including when
+    inference prefix caching is disabled.
     If False, we reuse KV cache from stale policies during generation
     (avoids recomputation at the cost of using slightly stale KV cache).
     """

tests/backends/skyrl_train/distributed/test_worker_dispatch.py

@@ -37,6 +37,43 @@ def _fft_dispatch_cfg(weight_sync_backend: str = "nccl") -> SimpleNamespace:
 class TestSaveWeights:
     """Tests for `WorkerDispatch.save_weights_for_sampler`"""
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("clear_kv_cache", [False, True])
+    @pytest.mark.parametrize("enable_prefix_caching", [False, True])
+    async def test_async_kv_offload_lifecycle(self, clear_kv_cache, enable_prefix_caching):
+        from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch
+
+        dispatch = WorkerDispatch.__new__(WorkerDispatch)
+        dispatch.colocate_all = False
+        dispatch.cfg = _fft_dispatch_cfg()
+        dispatch.cfg.trainer.fully_async = SimpleNamespace(enabled=True, clear_kv_cache_on_weight_sync=clear_kv_cache)
+        ie_cfg = dispatch.cfg.generator.inference_engine
+        ie_cfg.offload_kv_for_weight_sync = True
+        ie_cfg.enable_prefix_caching = enable_prefix_caching
+        dispatch._prepare_for_weight_sync = MagicMock()
+        dispatch._finish_weight_sync = MagicMock()
+        dispatch.ensure_active_adapter = MagicMock()
+        events = []
+        client = AsyncMock()
+        client.increment_weight_version = MagicMock()
+        client.pause_generation.side_effect = lambda: events.append("pause")
+        client.sleep_for_weight_sync.side_effect = lambda **kwargs: events.append(("sleep", kwargs["offload_kv"]))
+        client.wake_for_weight_sync.side_effect = lambda **kwargs: events.append(("wake", kwargs["tags"]))
+        client.resume_generation.side_effect = lambda: events.append("resume")
+        dispatch._inference_engine_client = client
+        dispatch._broadcast_to_inference_engines = MagicMock(side_effect=lambda *a, **kw: events.append("broadcast"))
+
+        await dispatch.save_weights_for_sampler()
+
+        assert events == [
+            "pause",
+            ("sleep", not clear_kv_cache),
+            ("wake", ["weights"]),
+            "broadcast",
+            ("wake", ["kv_cache"]),
+            "resume",
+        ]
+
     @pytest.mark.asyncio
     async def test_non_colocated_calls_pause_and_resume(self):
         from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch

tests/backends/skyrl_train/gpu/gpu_ci/test_offload_kv_weight_sync.py

@@ -3,8 +3,9 @@
 With ``generator.inference_engine.offload_kv_for_weight_sync=True`` (fully-async),
 ``WorkerDispatch.save_weights_for_sampler`` freezes in-flight requests (KEEP pause),
 offloads the KV cache to CPU, re-syncs weights into the freed space, then restores
-the KV cache. This asserts an in-flight request survives that offload/restore and
-finishes cleanly.
+the KV cache. With explicit KV clearing, it discards KV and resets running
+requests instead, including when prefix caching is disabled. In-flight requests
+must finish cleanly in either case.
 
 GPU Requirements: 2 GPUs (1 inference + 1 policy).
 
@@ -39,7 +40,7 @@
 LONG_PROMPT = "Tell me a very long, detailed story about a dragon who learns to code."
 
 
-def _offload_kv_cfg() -> SkyRLTrainConfig:
+def _offload_kv_cfg(clear_kv_cache: bool = False, enable_prefix_caching: bool = True) -> SkyRLTrainConfig:
     cfg = SkyRLTrainConfig()
     cfg.trainer.policy.model.path = MODEL
     cfg.trainer.critic.model.path = ""
@@ -48,11 +49,12 @@ def _offload_kv_cfg() -> SkyRLTrainConfig:
     cfg.trainer.placement.policy_num_gpus_per_node = 1
     cfg.trainer.remove_microbatch_padding = False
     cfg.trainer.logger = "console"
-    # Fully-async so the KV cache is offloaded to CPU and in-flight requests are
-    # frozen and resumed across the sync.
+    # Fully-async so requests remain in flight across the sync, either restoring
+    # their KV from CPU or recomputing it when explicit clearing is enabled.
     cfg.trainer.fully_async.enabled = True
-    cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = False
+    cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache
     ie = cfg.generator.inference_engine
+    ie.enable_prefix_caching = enable_prefix_caching
     ie.num_engines = 1
     ie.tensor_parallel_size = 1
     ie.run_engines_locally = True
@@ -90,7 +92,10 @@ def _sample_payload(prompt_token_ids: List[int], model: str, max_tokens: int) ->
 
 
 @pytest.mark.asyncio
-async def test_offload_kv_weight_sync_preserves_inflight_request(ray_init_fixture, monkeypatch):
+@pytest.mark.parametrize("clear_kv_cache,enable_prefix_caching", [(False, True), (True, False)])
+async def test_offload_kv_weight_sync_preserves_inflight_request(
+    ray_init_fixture, monkeypatch, clear_kv_cache, enable_prefix_caching
+):
     """An in-flight request must survive the KV offload/restore weight sync.
 
     While a long sample is mid-generation we run a real training step and
@@ -102,7 +107,7 @@ async def test_offload_kv_weight_sync_preserves_inflight_request(ray_init_fixtur
     # ignore_eos isn't in the default Tinker->vLLM forwarding map; widen it.
     monkeypatch.setitem(_ric._TINKER_SAMPLE_TO_VLLM_PARAM_MAP, "ignore_eos", "ignore_eos")
 
-    cfg = _offload_kv_cfg()
+    cfg = _offload_kv_cfg(clear_kv_cache, enable_prefix_caching)
     tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
     prompt_token_ids = _prompt_token_ids(tokenizer)
 
@@ -174,6 +179,7 @@ async def spy_wake(*args, **kwargs):
 
         # Confirm the KV-offloading path was actually exercised.
         assert sleep_calls, "sleep_for_weight_sync was never called — offload_kv path not taken"
+        assert sleep_calls[0][1]["offload_kv"] is not clear_kv_cache
         assert any(tags and "weights" in tags for tags in wake_tags), "weights were never woken for the broadcast"
         assert any(tags and "kv_cache" in tags for tags in wake_tags), "KV cache was never restored"
 

tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py

@@ -30,6 +30,7 @@ def _make_worker(worker_cls, handles_prefix_cache_reset: bool):
     worker._is_lora = False
     worker.cfg = SimpleNamespace(
         fully_async=SimpleNamespace(enabled=False, clear_kv_cache_on_weight_sync=False),
+        placement=SimpleNamespace(colocate_all=False),
         policy=SimpleNamespace(megatron_config=SimpleNamespace(lora_config=SimpleNamespace(merge_lora=False))),
     )
     worker._weight_transfer_sender = AsyncMock()
@@ -129,3 +130,27 @@ async def test_no_reset_when_prefix_caching_disabled(strategy, monkeypatch):
 
     client.reset_prefix_cache.assert_not_awaited()
     assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is False
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("strategy", ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)])
+@pytest.mark.parametrize("handles_reset", [False, True])
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+@pytest.mark.parametrize("clear_kv_cache", [False, True])
+async def test_async_running_request_reset(strategy, handles_reset, enable_prefix_caching, clear_kv_cache, monkeypatch):
+    _patch_collectives(monkeypatch)
+    worker = _make_worker(get_worker_cls(strategy), handles_prefix_cache_reset=handles_reset)
+    worker.cfg.fully_async.enabled = True
+    worker.cfg.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache
+    client = AsyncMock()
+    ie_cfg = SimpleNamespace(enable_prefix_caching=enable_prefix_caching, model_dtype="bfloat16")
+
+    await worker.broadcast_to_inference_engines(client, ie_cfg)
+
+    # The sender gets the same policy regardless of who owns the reset, and
+    # running requests require invalidation even without reusable prefix blocks.
+    assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is clear_kv_cache
+    if clear_kv_cache and not handles_reset:
+        client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True)
+    else:
+        client.reset_prefix_cache.assert_not_awaited()

tests/backends/skyrl_train/weight_sync/test_cache.py

@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache
+
+
+@pytest.mark.parametrize(
+    "fully_async,enable_prefix_caching,clear_kv_cache_on_weight_sync,expected",
+    [
+        (False, False, False, False),
+        (False, False, True, False),
+        (False, True, False, True),
+        (False, True, True, True),
+        (True, False, False, False),
+        (True, False, True, True),
+        (True, True, False, False),
+        (True, True, True, True),
+    ],
+)
+def test_weight_sync_cache_policy(fully_async, enable_prefix_caching, clear_kv_cache_on_weight_sync, expected):
+    assert (
+        should_reset_kv_cache(
+            enable_prefix_caching=enable_prefix_caching,
+            fully_async=fully_async,
+            clear_kv_cache_on_weight_sync=clear_kv_cache_on_weight_sync,
+        )
+        is expected
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+@pytest.mark.parametrize("clear_kv_cache", [False, True])
+async def test_sender_owned_async_reset(enable_prefix_caching, clear_kv_cache):
+    from skyrl.backends.skyrl_train.weight_sync.delta_strategy import (
+        DeltaWeightTransferSender,
+    )
+
+    client = AsyncMock()
+    sender = DeltaWeightTransferSender(SimpleNamespace(sync_dir="/unused"), client)
+    reset = should_reset_kv_cache(
+        enable_prefix_caching=enable_prefix_caching,
+        fully_async=True,
+        clear_kv_cache_on_weight_sync=clear_kv_cache,
+    )
+    await sender._apply_receiver_update({"target_version": 1}, rank=0, reset_prefix_cache=reset)
+
+    if clear_kv_cache:
+        client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True)
+        calls = [call[0] for call in client.mock_calls]
+        assert calls.index("pause_generation") < calls.index("reset_prefix_cache") < calls.index("start_weight_update")
+    else:
+        client.reset_prefix_cache.assert_not_awaited()
+    client.finish_weight_update.assert_awaited_once()
+    client.resume_generation.assert_awaited_once()

tests/train/test_config.py

@@ -420,13 +420,14 @@ def test_offload_kv_for_weight_sync_sync_trainer_ok():
 
 
 @pytest.mark.parametrize("clear_kv_cache", [False, True])
-def test_offload_kv_for_weight_sync_async_ok(clear_kv_cache):
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+def test_offload_kv_for_weight_sync_async_ok(clear_kv_cache, enable_prefix_caching):
     cfg = SkyRLTrainConfig()
     cfg.trainer.placement.colocate_all = False
     cfg.generator.inference_engine.offload_kv_for_weight_sync = True
     cfg.trainer.fully_async.enabled = True
     cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache
-    # Both clear_kv_cache settings are supported now.
+    cfg.generator.inference_engine.enable_prefix_caching = enable_prefix_caching
     validate_inference_engine_cfg(cfg)