configurable infinitstore server

vllm-project · Oct 10, 2024 · 0e6c551 · 0e6c551
1 parent dfe7474
commit 0e6c551
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vllm/distributed/kv_transfer/infinite.py b/vllm/distributed/kv_transfer/infinite.py
@@ -25,7 +25,8 @@ def __init__(self, model: str, tokens_per_page=16) -> None:
         self.tokens_per_page = tokens_per_page
         self.conn: InfinityConnection = InfinityConnection()
 
-        infinite_server = os.environ.get("PD_SEPARATE_STAGE", Default_Infinite_Server)
+        infinite_server = os.environ.get("INFINITE_STORE_SERVER", Default_Infinite_Server)
+        print("~~~~~~~~~~~~~connecting to infinite store server: ", infinite_server)
         self.conn.connect(infinite_server)
 
     def _compute_kv_cache_block_offsets(