[Core][Hybrid allocator + kv connector 1/n] Enable hybrid allocator + KV cache connector

tests/v1/core/test_scheduler.py

@@ -899,6 +899,7 @@ def test_kv_connector_basic():
     scheduler = create_scheduler(
         enable_prefix_caching=True,
         use_kv_connector=True,
+        disable_hybrid_kv_cache_manager=True,
     )
     NUM_TOTAL_BLOCKS = scheduler.kv_cache_manager.block_pool.get_num_free_blocks()
     BLOCK_SIZE = scheduler.cache_config.block_size
@@ -1028,6 +1029,7 @@ def test_kv_connector_unable_to_allocate():
         use_kv_connector=True,
         block_size=BLOCK_SIZE,
         num_blocks=NUM_BLOCKS,
+        disable_hybrid_kv_cache_manager=True,
     )
     NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
     scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
@@ -1111,6 +1113,7 @@ def test_kv_connector_handles_preemption():
         use_kv_connector=True,
         block_size=BLOCK_SIZE,
         num_blocks=NUM_BLOCKS,
+        disable_hybrid_kv_cache_manager=True,
     )
 
     NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
@@ -1327,6 +1330,7 @@ def create_scheduler_with_priority(
     block_size: int = 16,
     max_model_len: int | None = None,
     num_speculative_tokens: int | None = None,
+    disable_hybrid_kv_cache_manager: bool = False,
 ) -> Scheduler:
     """Create scheduler with priority policy enabled.
 
@@ -1351,6 +1355,7 @@ def create_scheduler_with_priority(
         disable_chunked_mm_input=disable_chunked_mm_input,
         enable_chunked_prefill=True,
         policy="priority",  # Enable priority scheduling
+        disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
     )
     model_config = ModelConfig(
         model=model,
@@ -1958,6 +1963,7 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv():
         num_blocks=5,  # Can hold 64 tokens (first block is null)
         block_size=16,  # Standard block size
         use_kv_connector=True,
+        disable_hybrid_kv_cache_manager=True,
     )
 
     # Create a request and schedule it

tests/v1/core/utils.py

@@ -46,6 +46,7 @@ def create_scheduler(
     num_speculative_tokens: int | None = None,
     skip_tokenizer_init: bool = False,
     async_scheduling: bool = False,
+    disable_hybrid_kv_cache_manager: bool = False,
 ) -> Scheduler | AsyncScheduler:
     """Create scheduler under test.
 
@@ -70,6 +71,7 @@ def create_scheduler(
         disable_chunked_mm_input=disable_chunked_mm_input,
         enable_chunked_prefill=True,
         async_scheduling=async_scheduling,
+        disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
     )
     model_config = ModelConfig(
         model=model,

tests/v1/kv_offload/test_cpu_offloading.py

@@ -27,6 +27,7 @@ def test_cpu_offloading(cpu_block_size: int) -> None:
         model="meta-llama/Llama-3.2-1B-Instruct",
         gpu_memory_utilization=0.5,
         kv_transfer_config=kv_transfer_config,
+        disable_hybrid_kv_cache_manager=True,
     )
 
     prompts = ["Hi " * 100]

vllm/config/vllm.py

@@ -40,11 +40,14 @@
     from transformers import PretrainedConfig
 
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
 else:
     PretrainedConfig = Any
 
     QuantizationConfig = Any
 
+    KVCacheConfig = Any
+
 logger = init_logger(__name__)
 
 
@@ -565,9 +568,6 @@ def __post_init__(self):
             if not current_platform.support_hybrid_kv_cache():
                 # Hybrid KV cache manager is not supported on non-GPU platforms.
                 self.scheduler_config.disable_hybrid_kv_cache_manager = True
-            if self.kv_transfer_config is not None:
-                # Hybrid KV cache manager is not compatible with KV transfer.
-                self.scheduler_config.disable_hybrid_kv_cache_manager = True
             if self.kv_events_config is not None:
                 # Hybrid KV cache manager is not compatible with KV events.
                 self.scheduler_config.disable_hybrid_kv_cache_manager = True

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -6,15 +6,18 @@
 from typing import TYPE_CHECKING, cast
 
 import vllm.envs as envs
+from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.base import (
     KVConnectorBase,
     KVConnectorBaseType,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    KVConnectorRole,
+    supports_hma,
+)
 from vllm.logger import init_logger
 
 if TYPE_CHECKING:
-    from vllm.config import VllmConfig
     from vllm.config.kv_transfer import KVTransferConfig
 
 logger = init_logger(__name__)
@@ -38,7 +41,7 @@ def loader() -> type[KVConnectorBase]:
     @classmethod
     def create_connector(
         cls,
-        config: "VllmConfig",
+        config: VllmConfig,
         role: KVConnectorRole,
     ) -> KVConnectorBase:
         if not envs.VLLM_USE_V1:
@@ -51,6 +54,15 @@ def create_connector(
         if kv_transfer_config is None:
             raise ValueError("kv_transfer_config must be set to create a connector")
         connector_cls = cls.get_connector_class(kv_transfer_config)
+
+        # check if the connector supports HMA
+        hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
+        if hma_enabled and not supports_hma(connector_cls):
+            raise ValueError(
+                f"Connector {connector_cls.__name__} does not support HMA but "
+                f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."
+            )
+
         logger.info(
             "Creating v1 connector with name: %s and engine_id: %s",
             connector_cls.__name__,

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

@@ -3,6 +3,13 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
     KVConnectorRole,
+    SupportsHMA,
+    supports_hma,
 )
 
-__all__ = ["KVConnectorRole", "KVConnectorBase_V1"]
+__all__ = [
+    "KVConnectorRole",
+    "KVConnectorBase_V1",
+    "supports_hma",
+    "SupportsHMA",
+]

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

@@ -70,6 +70,42 @@
 logger = init_logger(__name__)
 
 
+class SupportsHMA:
+    """
+    The class that indicates the corresponding connector supports hybrid memory
+    allocator (HMA).
+    This is required to use the connector together with hybrid memory allocator.
+    """
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called exactly once when a request has finished, before its blocks are
+        freed.
+
+        The connector may assumes responsibility for freeing the blocks
+        asynchronously by returning True.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        return False, None
+
+
+def supports_hma(connector: Any) -> bool:
+    if isinstance(connector, type):
+        return issubclass(connector, SupportsHMA)
+    else:
+        return isinstance(connector, SupportsHMA)
+
+
 class KVConnectorRole(enum.Enum):
     # Connector running in the scheduler process
     SCHEDULER = 0
@@ -370,7 +406,7 @@ def request_finished(
         Called exactly once when a request has finished, before its blocks are
         freed.
 
-        The connector may assumes responsibility for freeing the the blocks
+        The connector may assumes responsibility for freeing the blocks
         asynchronously by returning True.
 
         Returns:

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

@@ -10,6 +10,7 @@
     KVConnectorBase_V1,
     KVConnectorMetadata,
     KVConnectorRole,
+    SupportsHMA,
 )
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -23,8 +24,8 @@
 logger = init_logger(__name__)
 
 
-class LMCacheConnectorV1(KVConnectorBase_V1):
-    def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
+class LMCacheConnectorV1(SupportsHMA, KVConnectorBase_V1):
+    def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
         super().__init__(vllm_config=vllm_config, role=role)
         self._lmcache_engine = LMCacheConnectorV1Impl(vllm_config, role, self)
 
@@ -157,10 +158,10 @@ def build_connector_meta(
         """
         return self._lmcache_engine.build_connector_meta(scheduler_output)
 
-    def request_finished(
+    def request_finished(  # type: ignore[override]
         self,
         request: "Request",
-        block_ids: list[int],
+        block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
         """
         Called when a request has finished, before its blocks are freed.
@@ -171,5 +172,10 @@ def request_finished(
             get_finished().
             Optional KVTransferParams to be included in the request outputs
             returned by the engine.
+
+        Note:
+            This method intentionally uses tuple[list[int], ...] from
+            SupportsHMA interface instead of list[int] from
+            KVConnectorBase_V1 to support hybrid memory allocation.
         """
         return self._lmcache_engine.request_finished(request, block_ids)

vllm/v1/core/sched/scheduler.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
+import copy
 import itertools
 import time
 from collections import defaultdict
@@ -13,6 +13,7 @@
 from vllm.distributed.kv_transfer.kv_connector.v1 import (
     KVConnectorBase_V1,
     KVConnectorRole,
+    supports_hma,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
@@ -85,15 +86,14 @@ def __init__(
         # KV Connector pushes/pull of remote KVs for P/D and offloading.
         self.connector = None
         if self.vllm_config.kv_transfer_config is not None:
-            assert len(self.kv_cache_config.kv_cache_groups) == 1, (
-                "Multiple KV cache groups are not currently supported "
-                "with KV connectors"
-            )
             assert not self.is_encoder_decoder, (
                 "Encoder-decoder models are not currently supported with KV connectors"
             )
+
+            connector_vllm_config = copy.copy(self.vllm_config)
+            connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)
             self.connector = KVConnectorFactory.create_connector(
-                config=self.vllm_config, role=KVConnectorRole.SCHEDULER
+                config=connector_vllm_config, role=KVConnectorRole.SCHEDULER
             )
 
         self.kv_event_publisher = EventPublisherFactory.create(
@@ -1296,8 +1296,14 @@ def _connector_finished(
         if self.connector is None:
             return False, None
 
-        (block_ids,) = self.kv_cache_manager.get_block_ids(request.request_id)
-        return self.connector.request_finished(request, block_ids)
+        block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
+
+        if not supports_hma(self.connector):
+            # NOTE(Kuntai): We should remove this code path after we enforce
+            # all connectors to support HMA.
+            return self.connector.request_finished(request, block_ids[0])
+        else:
+            return self.connector.request_finished(request, block_ids)
 
     def _update_waiting_for_remote_kv(self, request: Request) -> bool:
         """

vllm/v1/worker/gpu_worker.py

@@ -327,6 +327,15 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
     def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
         """Allocate GPU KV cache with the specified kv_cache_config."""
 
+        # Init kv cache connector here, because it requires
+        # `kv_cache_config`.
+        # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
+        # because `initialize_kv_cache` will inject kv cache groups not
+        # related to kv cache connector (e.g. kv cache sharing layers).
+        connector_vllm_config = copy.copy(self.vllm_config)
+        connector_vllm_config.kv_cache_config = copy.copy(kv_cache_config)
+        ensure_kv_transfer_initialized(connector_vllm_config)
+
         if self.vllm_config.model_config.enable_sleep_mode:
             from vllm.device_allocator.cumem import CuMemAllocator
 
@@ -779,5 +788,3 @@ def init_worker_distributed_environment(
         parallel_config.pipeline_parallel_size,
         parallel_config.decode_context_parallel_size,
     )
-
-    ensure_kv_transfer_initialized(vllm_config)