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

skyrl/backends/skyrl_train/inference_engines/base.py

@@ -29,6 +29,11 @@ class InferenceEngineInput(TypedDict):
     sampling_params: Optional[Dict[str, Any]]
     session_ids: Optional[List[Hashable]]
     mm_features: Optional[List[MultiModalFeatures]]
+    # Optional prefix-cache salt applied to every prompt in this input. When set, it is forwarded to
+    # vLLM as the request ``cache_salt`` so that prefix-cache blocks are only shared between requests
+    # carrying the same salt. Typically the policy version (e.g. ``global_step``) -- see
+    # ``GeneratorConfig.use_cache_salt``.
+    cache_salt: Optional[str]
 
 
 class InferenceEngineOutput(TypedDict):

skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py

@@ -100,6 +100,7 @@ async def generate(
         prompt_token_ids = input_batch.get("prompt_token_ids")
         session_ids = input_batch.get("session_ids")
         sampling_params = input_batch.get("sampling_params")
+        cache_salt = input_batch.get("cache_salt")
 
         if (prompts is None and prompt_token_ids is None) or (prompts is not None and prompt_token_ids is not None):
             raise ValueError("Either `prompts` or `prompt_token_ids` must be provided, but not both.")
@@ -130,6 +131,7 @@ async def generate(
             engine_input = InferenceEngineInput(
                 prompt_token_ids=cur_prompt_token_ids,
                 sampling_params=sampling_params,
+                cache_salt=cache_salt,
             )
             tasks.append(asyncio.create_task(self.engines[engine_idx].generate(engine_input)))
             indices_list.append(prompt_ids)

skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py

@@ -180,6 +180,10 @@ async def generate(
                 payload = sampling_params.copy()
                 payload["model"] = self.model_name
                 payload["prompt"] = prompt_token_ids
+                # `cache_salt` is a top-level field on vLLM's CompletionRequest (not a sampling param),
+                # so we set it directly on the request body rather than inside `sampling_params`.
+                if (cache_salt := input_batch.get("cache_salt")) is not None:
+                    payload["cache_salt"] = cache_salt
                 request_url = f"{self.url}/v1/completions"
             else:
                 raise ValueError(f"Invalid engine backend: {self.engine_backend}")

skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py

@@ -144,7 +144,21 @@ def _preprocess_prompts(self, input_batch: InferenceEngineInput):
             SamplingParams(**request_sampling_params) if request_sampling_params is not None else SamplingParams()
         )
 
-        return prompt_token_ids, sampling_params
+        # `cache_salt` is a per-prompt field on vLLM's TokensPrompt (not a SamplingParams field), and is
+        # applied uniformly to every prompt in this input. Only effective when prefix caching is enabled.
+        cache_salt = input_batch.get("cache_salt")
+
+        return prompt_token_ids, sampling_params, cache_salt
+
+    @staticmethod
+    def _make_tokens_prompt(prompt_token_ids: List[int], cache_salt: Optional[str]) -> "TokensPrompt":
+        """Build a ``TokensPrompt``, attaching ``cache_salt`` only when it is set.
+
+        vLLM rejects an empty/``None`` ``cache_salt``, so we omit the field entirely when unset.
+        """
+        if cache_salt is not None:
+            return TokensPrompt(prompt_token_ids=prompt_token_ids, cache_salt=cache_salt)
+        return TokensPrompt(prompt_token_ids=prompt_token_ids)
 
     def _postprocess_outputs(self, outputs):
         """Common output processing logic."""
@@ -275,7 +289,7 @@ def _create_engine(self, *args, **kwargs):
         return vllm.LLM(*args, **kwargs)
 
     async def generate(self, input_batch: InferenceEngineInput) -> InferenceEngineOutput:
-        prompt_token_ids, sampling_params = self._preprocess_prompts(input_batch)
+        prompt_token_ids, sampling_params, cache_salt = self._preprocess_prompts(input_batch)
 
         # Check if LoRA is enabled and create LoRA requests
         lora_requests = None
@@ -291,7 +305,7 @@ async def generate(self, input_batch: InferenceEngineInput) -> InferenceEngineOu
 
         outputs = await asyncio.to_thread(
             self.llm.generate,
-            prompts=[TokensPrompt(prompt_token_ids=r) for r in prompt_token_ids],
+            prompts=[self._make_tokens_prompt(r, cache_salt) for r in prompt_token_ids],
             sampling_params=sampling_params,
             lora_request=lora_requests,
         )
@@ -507,7 +521,9 @@ async def _load_lora_from_disk(self, lora_path: str, lora_name: str = ""):
         result = await self.llm.add_lora(lora_request)
         return result
 
-    async def _collect_outputs(self, prompt_token_ids, request_id: str, sampling_params: SamplingParams):
+    async def _collect_outputs(
+        self, prompt_token_ids, request_id: str, sampling_params: SamplingParams, cache_salt: Optional[str] = None
+    ):
         """Collect outputs for a single prompt."""
         # Check if LoRA is enabled and create LoRA request
         final_output = None
@@ -523,7 +539,7 @@ async def _collect_outputs(self, prompt_token_ids, request_id: str, sampling_par
                 )
 
         async for request_output in self.llm.generate(
-            prompt=TokensPrompt(prompt_token_ids=prompt_token_ids),
+            prompt=self._make_tokens_prompt(prompt_token_ids, cache_salt),
             sampling_params=sampling_params,
             request_id=request_id,
             lora_request=lora_request,
@@ -534,14 +550,14 @@ async def _collect_outputs(self, prompt_token_ids, request_id: str, sampling_par
 
     async def generate(self, input_batch: InferenceEngineInput) -> InferenceEngineOutput:
         """Generate responses using vLLM's async engine."""
-        prompt_token_ids, sampling_params = self._preprocess_prompts(input_batch)
+        prompt_token_ids, sampling_params, cache_salt = self._preprocess_prompts(input_batch)
 
         tasks = []
         for prompt in prompt_token_ids:
             # Schedule the collection of outputs for each prompt.
             # Avoid duplicate request_ids
             request_id = str(uuid4().hex)
-            task = asyncio.create_task(self._collect_outputs(prompt, request_id, sampling_params))
+            task = asyncio.create_task(self._collect_outputs(prompt, request_id, sampling_params, cache_salt))
             tasks.append(task)
         outputs = await asyncio.gather(*tasks)
 

skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py

@@ -384,6 +384,7 @@ async def generate(
 
         session_ids = input_batch.get("session_ids")
         mm_features = input_batch.get("mm_features")
+        cache_salt = input_batch.get("cache_salt")
         get_logprobs = sampling_params.get("logprobs") is not None
 
         # Two semaphores decouple the generate and detokenize stages:
@@ -408,6 +409,7 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]:
                     session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None,
                     mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None,
                     model=model,
+                    cache_salt=cache_salt,
                 )
             async with gen_sem:
                 return await self._generate_single(
@@ -416,6 +418,7 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]:
                     session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None,
                     mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None,
                     model=model,
+                    cache_salt=cache_salt,
                 )
 
         async def _throttled_detokenize(token_ids: List[int]) -> str:
@@ -445,6 +448,7 @@ async def _generate_single(
         session_id: Optional[Any],
         model: str,
         mm_features: Optional[MultiModalFeatures] = None,
+        cache_salt: Optional[str] = None,
     ) -> Dict[str, Any]:
         """
         Generate completion for a single prompt.
@@ -469,6 +473,11 @@ async def _generate_single(
         }
         if mm_features:
             payload["features"] = mm_features
+        # `cache_salt` is a top-level request field (forwarded to vLLM's TokensPrompt), not a sampling
+        # param. The custom `/skyrl/v1/generate` endpoint honors it; the native data-plane endpoint
+        # forwards it when supported and otherwise ignores it.
+        if cache_salt is not None:
+            payload["cache_salt"] = cache_salt
 
         headers = {"Content-Type": "application/json"}
         if session_id:

skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py

@@ -437,9 +437,16 @@ async def _skyrl_generate(request: Request):
             body = await request.json()
             token_ids = body["token_ids"]
             sampling_params_dict = body.get("sampling_params", {})
+            cache_salt = body.get("cache_salt")
 
             sampling_params = VLLMSamplingParams(**sampling_params_dict)
-            prompt = TokensPrompt(prompt_token_ids=token_ids)
+            # `cache_salt` (when set) salts vLLM's prefix cache so blocks are only shared between
+            # requests carrying the same salt (e.g. the same policy version). vLLM rejects an empty
+            # salt, so we only attach it when present.
+            if cache_salt is not None:
+                prompt = TokensPrompt(prompt_token_ids=token_ids, cache_salt=cache_salt)
+            else:
+                prompt = TokensPrompt(prompt_token_ids=token_ids)
             request_id = random_uuid()
 
             final_res = None

skyrl/train/config/config.py

@@ -664,6 +664,16 @@ class GeneratorConfig(BaseConfig):
     eval_n_samples_per_prompt: int = 1
     zero_reward_on_non_stop: bool = False
     """Set reward to 0 when ``stop_reason`` is not ``"stop"`` (i.e., generation was truncated or aborted)."""
+    use_cache_salt: bool = False
+    """Salt the inference engine prefix cache with the policy version (``batch_metadata.global_step``).
+
+    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. Setting
+    this to ``True`` mixes a per-policy-version string (``cache_salt``) into the block hash so requests
+    at the same policy version share cache (throughput) while requests across versions do not
+    (correctness). This mirrors prime-rl's ``policy_version_at_start`` salt and is primarily useful for
+    fully-async RL, where the prefix cache is deliberately not fully cleared on every weight sync.
+    No-op unless prefix caching is enabled on the inference engine."""
     apply_overlong_filtering: bool = False
     """Apply DAPO Overlong Filtering: mask out all tokens in the loss mask for trajectories that
     exceed max length (truncated, no EOS token)."""

skyrl/train/generators/skyrl_gym_generator.py

@@ -265,6 +265,20 @@ def _post_process_agent_loop_output(
         """
         return agent_loop_output
 
+    def _compute_cache_salt(self, batch_metadata) -> Optional[str]:
+        """Derive a prefix-cache salt from the current policy version.
+
+        When ``generator.use_cache_salt`` is enabled, returns a string that is stable within a single
+        ``generate`` batch (one policy version) but changes across weight updates, so vLLM's prefix cache
+        is shared within a version and isolated across versions. The policy model name is included so
+        distinct LoRA adapters / tenants do not collide. Returns ``None`` (no salting) when disabled or
+        when no ``batch_metadata`` is available (e.g. ad-hoc ``generate`` calls without a global step).
+        """
+        if not self.generator_cfg.use_cache_salt or batch_metadata is None:
+            return None
+        version = f"{self.policy_model_name}@" if self.policy_model_name is not None else ""
+        return f"{version}{batch_metadata.global_step}"
+
     async def agent_loop(
         self,
         prompt: ConversationType,
@@ -274,6 +288,7 @@ async def agent_loop(
         max_input_length: int,
         sampling_params: Optional[Dict[str, Any]] = None,
         trajectory_id: Optional[TrajectoryID] = None,
+        cache_salt: Optional[str] = None,
     ) -> Union[TrajectoryOutput, StepWiseOutput]:
         """
         Multi-turn generation loop that executes a single trajectory.
@@ -391,6 +406,7 @@ async def agent_loop(
                     prompt_token_ids=[agent_loop_state.input_ids],
                     session_ids=[session_id],
                     sampling_params=sampling_params,
+                    cache_salt=cache_salt,
                 )
                 engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)
                 output = engine_output["responses"][0]
@@ -707,6 +723,7 @@ async def generate_batched(
         env_extras: List[Dict[str, Any]],
         max_tokens: int,
         sampling_params: Optional[Dict[str, Any]] = None,
+        cache_salt: Optional[str] = None,
     ) -> GeneratorOutput:
         """
         Single-turn batched generation (can use the synchronous offline engine)
@@ -738,7 +755,9 @@ async def generate_batched(
             tokenize=True,
             return_dict=False,
         )
-        engine_input = InferenceEngineInput(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params)
+        engine_input = InferenceEngineInput(
+            prompt_token_ids=prompt_token_ids, sampling_params=sampling_params, cache_salt=cache_salt
+        )
         engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)
         outputs = engine_output["responses"]
         responses = engine_output["response_ids"]
@@ -816,8 +835,14 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         max_tokens = self.generator_cfg.sampling_params.max_generate_length
         max_input_length = self.generator_cfg.max_input_length
 
+        # Prefix-cache salt derived from the policy version. Captured once per `generate` call so every
+        # trajectory in this batch shares one salt (matches prime-rl's "version at start of group").
+        cache_salt = self._compute_cache_salt(input_batch.get("batch_metadata", None))
+
         if self.batched:
-            return await self.generate_batched(prompts, env_classes, env_extras, max_tokens, sampling_params)
+            return await self.generate_batched(
+                prompts, env_classes, env_extras, max_tokens, sampling_params, cache_salt=cache_salt
+            )
 
         # Async agent loop to generate trajectories in parallel.
         tasks = []
@@ -831,6 +856,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                     max_input_length,
                     sampling_params=sampling_params,
                     trajectory_id=trajectory_ids[i] if trajectory_ids is not None else None,
+                    cache_salt=cache_salt,
                 )
             )
 

skyrl/train/generators/skyrl_vlm_generator.py

@@ -74,6 +74,7 @@ async def agent_loop(
         max_input_length: int,
         sampling_params: Optional[Dict[str, Any]] = None,
         trajectory_id: Optional[TrajectoryID] = None,
+        cache_salt: Optional[str] = None,
     ) -> TrajectoryOutput:
         """Multi-turn VLM generation loop for a single trajectory.
         The conversation is treated as the source of truth and re-tokenized each step.
@@ -155,6 +156,7 @@ async def agent_loop(
                     session_ids=[session_id],
                     sampling_params=current_sampling_params,
                     mm_features=[latest_features] if latest_features is not None else None,
+                    cache_salt=cache_salt,
                 )
                 engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)