From 0e6c551fc7a1bea686ef1200d0d2a6fcfae6ab6d Mon Sep 17 00:00:00 2001 From: chenqianfzh Date: Thu, 10 Oct 2024 22:33:41 +0000 Subject: [PATCH] configurable infinitstore server --- vllm/distributed/kv_transfer/infinite.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/infinite.py b/vllm/distributed/kv_transfer/infinite.py index 5334b72a65bec..fac6b3de9fabf 100644 --- 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(