[async] add flag to toggle whether to clear kv cache after weight sync for fully async training, update default max_staleness_steps to 0

[erictang000]

Previously, we always reset prefix cache prior to weight sync, even for fully async training:

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

Line 1274 in 649b16f

if use_prefix_cache and torch.distributed.get_rank() == 0:

Adds a new flag trainer.fully_async.clear_kv_cache_on_weight_sync which defaults to False, meaning the default behavior during fully async training is to keep stale KV cache.

Also updates the default for trainer.fully_async.max_staleness_steps from 4 -> 0, to be able to determine whether or not fully async training is enabled. All existing fully async scripts already manually set max_staleness_steps, and users should manually set this to a non zero value for all fully async scripts anyway.

[gemini-code-assist]

Code Review

This pull request introduces a new configuration option clear_kv_cache_on_weight_sync to FullyAsyncConfig and updates both FSDP and Megatron worker backends to skip resetting the KV cache during fully async training when this option is disabled. The review feedback suggests correcting the docstring for max_staleness_steps to state it should be set > 0 instead of > 1 for fully async training. Additionally, the reviewer recommends encapsulating the skip_kv_cache_reset logic as a property on FullyAsyncConfig to avoid code duplication across the worker backends.

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.

[erictang000]

verified similar performance with not clearing kv cache

https://api.wandb.ai/links/sky-posttraining-uc-berkeley/p8c50yl8

[CharlieFRuan]

I think max_staleness = 0 with fully async might still be used for debugging. Treating > 0 as a flag for fully async is a bit implicit to me.

Could we set clear_kv_cache_on_weight_sync default to True instead (which applies to sync RL), and add it as a flag to be false for the example fully async scripts?