initial noodling

Signed-off-by: Will Eaton <weaton@redhat.com>

Author: wseaton

.gitignore

@@ -202,3 +202,4 @@ shellcheck*/
 
 # Ingore moe/marlin_moe gen code
 csrc/moe/marlin_moe_wna16/kernel_*
+uv.lock

vllm/config.py

@@ -10,7 +10,7 @@
 import textwrap
 import uuid
 import warnings
-from collections import Counter
+from collections import Counter, defaultdict
 from contextlib import contextmanager
 from dataclasses import (MISSING, Field, asdict, field, fields, is_dataclass,
                          replace)
@@ -33,6 +33,8 @@
 import vllm.envs as envs
 from vllm import version
 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorHandshakeMetadata)
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization import (QUANTIZATION_METHODS,
                                                      QuantizationMethods,
@@ -1511,6 +1513,12 @@ class CacheConfig:
     num_cpu_blocks: Optional[int] = field(default=None, init=False)
     """The number of blocks to allocate for CPU memory."""
 
+    transfer_handshake_metadata: dict[int, dict[int, 
+        KVConnectorHandshakeMetadata]] = field(
+        default_factory=lambda: defaultdict(dict), 
+        init=False)
+    """Metadata for the KV connector handshake."""
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,

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

@@ -32,9 +32,11 @@
 
 import enum
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import torch
+import msgspec
+from pydantic_core import core_schema
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -62,18 +64,58 @@ class KVConnectorMetadata:
     Abstract Metadata used to communicate between the
     Scheduler KVConnector and Worker KVConnector.
     """
-    pass
+    
+    def __init__(self):
+        pass
+
 
+class KVConnectorHandshakeMetadata(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        # required for @cached_property.
+        dict=True):
+    """
+    Metadata optionally used for out of band connector handshake between P/D workers.
+    """
+    connector_type: str = "base"
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, 
+        _source_type: Any, 
+        _handler: Callable[[Any], core_schema.CoreSchema]
+    ) -> core_schema.CoreSchema:
+        """bridge msgspec.Struct with pydantic for schema generation"""
+        return core_schema.no_info_after_validator_function(
+            cls,
+            core_schema.dict_schema()
+        )
+
+class KVConnectorTransferMetadata(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        dict=True):
+    """
+    Wrapper for transfer handshake metadata sent between engine and utils.
+    """
+    tensor_parallel_rank: int
+    data_parallel_rank: int
+    content: Optional[dict]
+    
 
 class KVConnectorBase_V1(ABC):
 
-    def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
+    def __init__(self,
+                 vllm_config: "VllmConfig",
+                 role: KVConnectorRole):
         logger.warning(
             "Initializing KVConnectorBase_V1. This API is experimental and "
             "subject to change in the future as we iterate the design.")
         self._connector_metadata = KVConnectorMetadata()
         self._vllm_config = vllm_config
         self._role = role
+        self._handshake_metadata: Optional[KVConnectorHandshakeMetadata] = None
+
 
     @property
     def role(self) -> KVConnectorRole:
@@ -104,7 +146,7 @@ def clear_connector_metadata(self) -> None:
         """
         self._connector_metadata = KVConnectorMetadata()
 
-    def _get_connector_metadata(self) -> KVConnectorMetadata:
+    def get_connector_metadata(self) -> KVConnectorMetadata:
         """Get the connector metadata.
 
         This function should only be called inside the connector.
@@ -201,6 +243,31 @@ def get_finished(
         """
         return None, None
 
+    def set_handshake_metadata(
+        self, handshake_metadata: KVConnectorHandshakeMetadata) -> None:
+        """
+        Set the handshake metadata for the connector.
+
+        This metadata is used for out-of-band connector handshake
+        between P/D workers.
+        
+        Args:
+            handshake_metadata (KVConnectorHandshakeMetadata): the handshake
+                metadata.
+        """
+        self._handshake_metadata = handshake_metadata
+
+
+    def get_handshake_metadata(
+        self) -> Optional[KVConnectorHandshakeMetadata]:
+        """
+        Get the handshake metadata for the connector.
+
+        Returns:
+            KVConnectorHandshakeMetadata: the handshake metadata.
+        """
+        return self._handshake_metadata
+
     # ==============================
     # Scheduler-side methods
     # ==============================