Skip to content

[generator] Add per global_step cache salt for each trajectory to invalidate super stale KV#1836

Open
erictang000 wants to merge 8 commits into
NovaSky-AI:mainfrom
erictang000:cache_salt
Open

[generator] Add per global_step cache salt for each trajectory to invalidate super stale KV#1836
erictang000 wants to merge 8 commits into
NovaSky-AI:mainfrom
erictang000:cache_salt

Conversation

@erictang000

@erictang000 erictang000 commented Jun 26, 2026

Copy link
Copy Markdown
Collaborator

Overview

Adds generator.use_cache_salt config - defaults to True.

vLLM hashes prefix-cache blocks by token content alone, so two requests with an identical token prefix share KV blocks even across a weight update -- where the cached KV is now stale.

When generator.use_cache_salt=True (the default), a per-policy-version string (cache_salt) is mixed into the block hash so trajectories that begin at the same policy version share cache (throughput) while requests across versions do not. The salt is keyed on the inference engine's weight version (the number of weight
syncs broadcast to the engines), captured at the start of each generate batch (for a multi-turn trajectory, we still get prefix hits across turns, since the salt stays consistent throughout the lifetime of the trajectory).

This matters for fully-async RL, where the prefix cache is deliberately not fully cleared on every weight sync. It is effectively a no-op for synchronous training (which already resets the prefix cache on each weight sync) and whenever prefix caching is disabled on the inference engine, so it is safe to leave on by default

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a prefix-cache salting mechanism (cache_salt) derived from the policy version to prevent stale KV cache sharing across weight updates in vLLM. The changes propagate this salt across inference engines, clients, servers, and generators. The review feedback points out a potential AttributeError in _compute_cache_salt if batch_metadata is a dictionary, and recommends checking for truthiness (if cache_salt:) instead of is not None across multiple files to prevent empty strings from being forwarded to vLLM, which rejects them.

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.

Comment thread skyrl/train/generators/skyrl_gym_generator.py Outdated
Comment thread skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py Outdated
Comment thread skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py
Comment thread skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py Outdated
Comment thread skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py Outdated
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.

2 participants