A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 19:00 UTC
[Feat][KVConnector] Add `bind_gpu_block_pool()` to KVConnectorBase_V1 (#39654)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Yifan Qiao committed
May 13, 2026 at 02:10 UTC
13bf2421009a001b79751666695623ad8b9f29b2
5 files changed
+24
-9
tests/v1/kv_connector/unit/test_multi_connector.py
+7
-5
@@ -261,10 +261,11 @@ def test_multi_example_connector_consistency():
261
)
262
263
events = get_connector_events()
264
- # First event is set_xfer_handshake_metadata from initialization, then
265
- # on_new_request when the request is enqueued, then
266
- # get_num_new_matched_tokens and update_state_after_alloc from generate().
267
- assert events["storage1-SCHEDULER"][:5] == [
264
+ # First event is bind_gpu_block_pool from initialization, then
265
+ # set_xfer_handshake_metadata, then on_new_request when the request is enqueued,
266
+ # then get_num_new_matched_tokens and update_state_after_alloc from generate().
267
+ assert events["storage1-SCHEDULER"][:6] == [
268
+ "bind_gpu_block_pool",
269
"set_xfer_handshake_metadata",
270
"on_new_request",
271
"get_num_new_matched_tokens 0",
@@ -283,7 +284,8 @@ def test_multi_example_connector_consistency():
284
"wait_for_layer_load",
285
"save_kv_layer",
286
]
286
- assert events["storage2-SCHEDULER"][:5] == [
287
+ assert events["storage2-SCHEDULER"][:6] == [
288
+ "bind_gpu_block_pool",
289
"set_xfer_handshake_metadata",
290
"on_new_request",
291
"get_num_new_matched_tokens 0",
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+11
@@ -62,6 +62,7 @@ if TYPE_CHECKING:
62
PromMetricT,
63
)
64
from vllm.forward_context import ForwardContext
65
+ from vllm.v1.core.block_pool import BlockPool
66
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
67
from vllm.v1.kv_cache_interface import KVCacheConfig
68
from vllm.v1.request import Request
@@ -439,6 +440,16 @@ class KVConnectorBase_V1(ABC):
440
# Scheduler-side methods
441
# ==============================
442
443
+ def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
444
+ """
445
+ Bind the GPU block pool to the connector for per-GPU block status tracking.
446
+ For example, inc/dec ref counts, or iterate over the prefix cache blocks.
447
+
448
+ Args:
449
+ gpu_block_pool: the GPU block pool.
450
+ """
451
+ return
452
+
453
@abstractmethod
454
def get_num_new_matched_tokens(
455
self,
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+5
@@ -35,6 +35,7 @@ from vllm.v1.outputs import KVConnectorOutput
35
if TYPE_CHECKING:
36
from vllm.distributed.kv_events import KVCacheEvent
37
from vllm.forward_context import ForwardContext
38
+ from vllm.v1.core.block_pool import BlockPool
39
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
40
from vllm.v1.kv_cache_interface import KVCacheConfig
41
from vllm.v1.request import Request
@@ -227,6 +228,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
228
for c in self._connectors:
229
c.register_kv_caches(kv_caches)
230
231
+ def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
232
+ for c in self._connectors:
233
+ c.bind_gpu_block_pool(gpu_block_pool)
234
+
235
# We must override the base class method here because we need to bind
236
# the metadata to each connector in the order of the connectors in the
237
# MultiKVConnectorMetadata.
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
-1
@@ -165,7 +165,6 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
165
166
# --- Scheduler-side methods ---
167
168
- # NOTE: New API only for SimpleCPUOffloadConnector.
168
def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
169
if self.scheduler_manager is not None:
170
self.scheduler_manager.bind_gpu_block_pool(gpu_block_pool)
vllm/v1/core/sched/scheduler.py
+1
-3
@@ -235,9 +235,7 @@ class Scheduler(SchedulerInterface):
235
)
236
# Bind GPU block pool to the KV connector. This must happen after
237
# kv_cache_manager is constructed so block_pool is available.
238
- if self.connector is not None and hasattr(
239
- self.connector, "bind_gpu_block_pool"
240
- ):
238
+ if self.connector is not None:
239
self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)
240
241
self.use_pp = self.parallel_config.pipeline_parallel_size > 1