Skip to content

fix(inference): reset requests when discarding KV cache - #2147

Draft
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-async-kv-reset
Draft

fix(inference): reset requests when discarding KV cache#2147
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-async-kv-reset

Conversation

@bvolpato

@bvolpato bvolpato commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Fully async weight sync can discard the KV cache of KEEP-paused requests without resetting those requests. With KV offloading and clear_kv_cache_on_weight_sync=true, dispatch skips the CPU backup, but the FSDP and Megatron workers previously skipped request reset when prefix caching was disabled. Generation then resumed with computed-token/block state referring to discarded KV contents.

Both workers now honor the explicit async clear setting independently of reusable prefix caching. They share the same reset policy and retain sender-owned reset handling. Synchronous training without prefix caching still skips reset.

Reproduction and validation

The affected configuration is accepted today:

trainer.placement.colocate_all: false
trainer.fully_async.enabled: true
trainer.fully_async.clear_kv_cache_on_weight_sync: true
generator.inference_engine.offload_kv_for_weight_sync: true
generator.inference_engine.enable_prefix_caching: false
uv run --isolated --extra skyrl-train --extra dev pytest \
  tests/backends/skyrl_train/weight_sync/test_cache.py \
  tests/backends/skyrl_train/distributed/test_worker_dispatch.py \
  tests/train/test_config.py -k 'cache or offload or TestSaveWeights' -q

uv run --isolated --extra fsdp --extra dev pytest \
  tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py \
  -m 'not megatron' -q
  • 32 selected CPU tests passed for policy, sender ownership, dispatch, and configuration.
  • The actual original FSDP worker failed both new disabled-prefix/clear-enabled cases; 9 controls passed. The fixed worker passed all 11 cases. This local run stubbed unused FlashAttention imports and mocked GPU collectives, as the existing worker tests do.
  • Added the corresponding Megatron integration cases and extended the GPU offload smoke test with clear=true, prefix=false. These have not run on this CPU-only host; the draft still needs GPU validation. The smoke test does not establish numerical parity.
  • Ruff, Black, and gitleaks hooks passed.
  • GitHub GPU CI stopped during Anyscale setup because the fork job has no ANYSCALE_CLI_TOKEN; it did not execute tests. Upstream code-quality and gym checks passed; remaining CPU jobs were still running when submitted.

Pinned vLLM's request reset preempts running requests independently of prefix caching. Related to the offload path introduced in #1934.

Downsides

With async clearing explicitly enabled and prefix caching disabled, active requests now recompute their KV after each sync. This adds prefill work proportional to their retained context, which is the intended cost of discarding KV. The default clear_kv_cache_on_weight_sync=false path is unchanged.

Risk and rollback

The change affects FSDP and Megatron request-reset decisions during weight sync. Check generation completion and sync latency in the added GPU case before marking ready. To roll back safely, disable generator.inference_engine.offload_kv_for_weight_sync before reverting so in-flight KV is retained rather than discarded.

Signed-off-by: bvolpato <brunocvcunha@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant