first nixl version

Author: bukeao

vllm/distributed/kv_transfer/kv_connector/v1/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-    KVConnectorBase_V1, KVConnectorRole)
+    KVConnectorBase_V1, KVConnectorRole, KVTransferParams)
 
-__all__ = ["KVConnectorRole", "KVConnectorBase_V1"]
+__all__ = ["KVConnectorRole", "KVConnectorBase_V1", "KVTransferParams"]

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -56,6 +56,21 @@ class KVConnectorRole(enum.Enum):
     # Connector running in the worker process
     WORKER = 1
 
+class KVTransferParams:
+    """
+    Abstract KVTransferParams used to send KVTransfer
+    parameters between instances of vLLM.
+    
+    Specific instances of KVConnector customize this
+    method for serializing / deserializing msgs sent
+    via the HTTP protocol.
+    """
+
+    @staticmethod
+    def from_raw_dict(
+            raw_dict: Optional[dict[str,
+                                    Any]]) -> Optional["KVTransferParams"]:
+        return None
 
 class KVConnectorMetadata:
     """
@@ -66,6 +81,7 @@ class KVConnectorMetadata:
 
 
 class KVConnectorBase_V1(ABC):
+    _KVTransferParams = KVTransferParams
 
     def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         logger.warning(
@@ -204,7 +220,13 @@ def get_finished(
     # ==============================
     # Scheduler-side methods
     # ==============================
-
+    def set_kv_transfer_params(self, request: "Request"):
+        """Parse raw KV Transfer params."""
+        assert request.kv_transfer_params is None
+        kv_transfer_params = self._KVTransferParams.from_raw_dict(
+            request.raw_kv_transfer_params)
+        request.kv_transfer_params = kv_transfer_params
+        
     @abstractmethod
     def get_num_new_matched_tokens(
         self,