[Tunix] Clear and reinitialize vLLM KV cache when loading weights.

PiperOrigin-RevId: 875844163

Author: wang2yn84

tunix/generate/vllm_sampler.py

@@ -159,6 +159,12 @@ def update_params(
       filter_types: Optional[Tuple[Any, ...]] = None,
   ):
     del filter_types
+    if self.llm is not None:
+      self.llm.reset_prefix_cache()
+      self.llm.collective_rpc("delete_kv_cache") # will free hbm
+    elif self._driver is not None:
+      self._driver.llm_engine.reset_prefix_cache()
+      self._driver.llm_engine.collective_rpc("delete_kv_cache")
 
     if self.to_hf_key_mappings:
       # Mapped Weight Sync (e.g. Vanilla -> vLLM)
@@ -192,6 +198,11 @@ def update_params(
           reshard_fn=reshard.reshard_pytree,
       )
 
+    if self.llm is not None:
+      self.llm.collective_rpc("reinitialize_kv_cache")
+    elif self._driver is not None:
+      self._driver.llm_engine.collective_rpc("reinitialize_kv_cache")
+
   def load_checkpoint(self, path_or_weights: str | jaxtyping.PyTree):
     # TODO(b/434741253): Consider support orbax checkpoint loading
     if isinstance(path_or_weights, jaxtyping.PyTree):