@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 647464719b
Fork 0

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 12, 21:00 UTC

[KVConnector][Core] Support cross-layer KV blocks (#27743)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Nov 20, 2025 at 20:09 UTC 647464719b131963dccdc3a28cfe52d1af293cda
15 files changed +453 -90
tests/v1/kv_connector/unit/test_offloading_connector.py
+6 -2
@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
19 )
20 from vllm.forward_context import ForwardContext
21 from vllm.utils.hashing import sha256
22 +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
23 from vllm.v1.core.kv_cache_utils import (
24 BlockHash,
25 get_request_block_hasher,
@@ -92,7 +93,7 @@ class MockOffloadingSpec(OffloadingSpec):
93 return self.manager
94
95 def get_handlers(
95 - self, _
96 + self, _, __
97 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
98 yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
99 yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
@@ -138,7 +139,10 @@ class RequestRunner:
139 self.worker_connector = OffloadingConnector(vllm_config, KVConnectorRole.WORKER)
140
141 # register worker kv_caches to enable OffloadingWorker creations
141 - self.worker_connector.register_kv_caches(kv_caches={"a": torch.empty(0)})
142 + self.worker_connector.register_cross_layers_kv_cache(
143 + kv_cache=torch.empty(0),
144 + attn_backend=FlashAttentionBackend,
145 + )
146
147 # extract connector of scheduler
148 scheduler_connector = self.scheduler.connector
tests/v1/kv_offload/test_cpu_offloading.py
+93 -52
@@ -12,8 +12,10 @@ from tqdm import tqdm
12 from vllm import LLM, SamplingParams, TokensPrompt
13 from vllm.config import KVEventsConfig, KVTransferConfig
14 from vllm.distributed.kv_events import BlockStored, KVEventBatch
15 +from vllm.utils.system_utils import set_env_var
16
16 -CPU_BLOCK_SIZES = [16, 48]
17 +CPU_BLOCK_SIZES = [48]
18 +ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER"]
19
20
21 class MockSubscriber:
@@ -63,8 +65,88 @@ class MockSubscriber:
65 self.sub.close()
66
67
68 +def _latency_test(llm: LLM, subscriber: MockSubscriber):
69 + sampling_params = SamplingParams(max_tokens=1)
70 +
71 + num_times_cpu_better_than_cold = 0
72 + num_tests = 10
73 + total_cold_time = 0.0
74 + total_gpu_hit_time = 0.0
75 + total_cpu_hit_time = 0.0
76 + prompt_token_ids = [0] * 10001
77 + for i in tqdm(range(num_tests), desc="Running tests"):
78 + prompt_token_ids[0] = i
79 + prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
80 +
81 + # run generation - this should trigger saving KV cache
82 + start_time = time.time()
83 + llm.generate(prompts, sampling_params, use_tqdm=False)
84 + cold_time = time.time() - start_time
85 + total_cold_time += cold_time
86 +
87 + # run generation again - should hit the GPU prefix cache
88 + start_time = time.time()
89 + llm.generate(prompts, sampling_params, use_tqdm=False)
90 + gpu_hit_time = time.time() - start_time
91 + total_gpu_hit_time += gpu_hit_time
92 +
93 + # reset prefix cache to avoid GPU hit.
94 + llm.reset_prefix_cache()
95 +
96 + assert subscriber.get_new_cpu_stored_events()
97 +
98 + # run generation again - this should trigger loading from CPU
99 + start_time = time.time()
100 + llm.generate(prompts, sampling_params, use_tqdm=False)
101 + cpu_hit_time = time.time() - start_time
102 + total_cpu_hit_time += cpu_hit_time
103 +
104 + if cpu_hit_time < cold_time:
105 + num_times_cpu_better_than_cold += 1
106 +
107 + print("Average times:")
108 + print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
109 + print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms")
110 + print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms")
111 +
112 + assert num_times_cpu_better_than_cold >= 0.8 * num_tests
113 +
114 +
115 +def _accuracy_test(llm: LLM, subscriber: MockSubscriber):
116 + sampling_params = SamplingParams(max_tokens=1)
117 + cpu_block_size = (
118 + llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config[
119 + "block_size"
120 + ]
121 + )
122 +
123 + subscriber.get_new_cpu_stored_events()
124 +
125 + # prepend prompt to be cpu block aligned
126 + prompt = "Let's count to 10. One, two, three, four,"
127 + while (
128 + len(llm.generate(prompt, use_tqdm=False)[0].prompt_token_ids) % cpu_block_size
129 + != 0
130 + ):
131 + prompt = ". " + prompt
132 +
133 + assert subscriber.get_new_cpu_stored_events()
134 +
135 + test_count = 100
136 + success_count = 0
137 + for i in range(test_count):
138 + if (
139 + llm.generate(prompt, sampling_params, use_tqdm=False)[0].outputs[0].text
140 + == " five"
141 + ):
142 + success_count += 1
143 +
144 + assert success_count >= 0.5 * test_count
145 +
146 +
147 @pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
67 -def test_cpu_offloading(cpu_block_size: int) -> None:
148 +@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
149 +def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
150 """
151 Tests OffloadingConnector with CPUOffloadingSpec.
152 """
@@ -92,61 +174,20 @@ def test_cpu_offloading(cpu_block_size: int) -> None:
174 topic="test",
175 )
176
95 - llm = LLM(
96 - model="meta-llama/Llama-3.2-1B-Instruct",
97 - gpu_memory_utilization=0.5,
98 - kv_events_config=kv_events_config,
99 - kv_transfer_config=kv_transfer_config,
100 - )
101 -
102 - sampling_params = SamplingParams(temperature=0, max_tokens=1)
177 + with set_env_var("VLLM_ATTENTION_BACKEND", attn_backend):
178 + llm = LLM(
179 + model="meta-llama/Llama-3.2-1B-Instruct",
180 + gpu_memory_utilization=0.5,
181 + kv_events_config=kv_events_config,
182 + kv_transfer_config=kv_transfer_config,
183 + )
184
185 events_endpoint = events_endpoint.replace("*", "127.0.0.1")
186 subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
187
188 try:
108 - num_times_cpu_better_than_cold = 0
109 - num_tests = 10
110 - total_cold_time = 0.0
111 - total_gpu_hit_time = 0.0
112 - total_cpu_hit_time = 0.0
113 - prompt_token_ids = [0] * 10001
114 - for i in tqdm(range(num_tests), desc="Running tests"):
115 - prompt_token_ids[0] = i
116 - prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
117 -
118 - # run generation - this should trigger saving KV cache
119 - start_time = time.time()
120 - llm.generate(prompts, sampling_params, use_tqdm=False)
121 - cold_time = time.time() - start_time
122 - total_cold_time += cold_time
123 -
124 - # run generation again - should hit the GPU prefix cache
125 - start_time = time.time()
126 - llm.generate(prompts, sampling_params, use_tqdm=False)
127 - gpu_hit_time = time.time() - start_time
128 - total_gpu_hit_time += gpu_hit_time
129 -
130 - # reset prefix cache to avoid GPU hit.
131 - llm.reset_prefix_cache()
132 -
133 - assert subscriber.get_new_cpu_stored_events()
134 -
135 - # run generation again - this should trigger loading from CPU
136 - start_time = time.time()
137 - llm.generate(prompts, sampling_params, use_tqdm=False)
138 - cpu_hit_time = time.time() - start_time
139 - total_cpu_hit_time += cpu_hit_time
140 -
141 - if cpu_hit_time < cold_time:
142 - num_times_cpu_better_than_cold += 1
143 -
144 - print("Average times:")
145 - print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
146 - print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms")
147 - print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms")
148 -
149 - assert num_times_cpu_better_than_cold >= 0.8 * num_tests
189 + _latency_test(llm, subscriber)
190 + _accuracy_test(llm, subscriber)
191 finally:
192 subscriber.close()
193 del llm
tests/v1/worker/test_gpu_model_runner.py
+4 -1
@@ -483,7 +483,10 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
483 # Permutation that gets you back to expected kv shape
484 for test_stride in ((1, 4, 0, 2, 3), (0, 1, 2, 3, 4)):
485
486 - def rnd_stride_order(test_stride=test_stride):
486 + def rnd_stride_order(
487 + include_num_layers_dimension: bool = False, test_stride=test_stride
488 + ):
489 + assert not include_num_layers_dimension
490 return test_stride
491
492 # Patch the attention backend class and re-trigger the KV cache creation
vllm/attention/backends/abstract.py
+28 -1
@@ -76,7 +76,34 @@ class AttentionBackend(ABC):
76 raise NotImplementedError
77
78 @staticmethod
79 - def get_kv_cache_stride_order() -> tuple[int, ...]:
79 + def get_kv_cache_stride_order(
80 + include_num_layers_dimension: bool = False,
81 + ) -> tuple[int, ...]:
82 + """
83 + Get the physical (memory layout) ordering of the kv cache dimensions.
84 + e.g. if the KV cache shape is
85 + [2, num_blocks, block_size, num_heads, head_size],
86 + and get_kv_cache_stride_order returns (1, 3, 0, 2, 4) then the physical
87 + ordering of dimensions is
88 + [num_blocks, num_heads, 2, block_size, head_size].
89 +
90 + If this function is unimplemented / raises NotImplementedError,
91 + the physical layout of the KV cache will match the logical shape.
92 +
93 + Args:
94 + include_num_layers_dimension: if True, includes an additional
95 + num_layers dimension, which is assumed to be prepended
96 + to the logical KV cache shape.
97 + With the above example, a return value (2, 4, 0, 1, 3, 5)
98 + corresponds to
99 + [num_blocks, num_heads, num_layers, 2, block_size, head_size].
100 +
101 + If an additional dimension is NOT included in the returned
102 + tuple, the physical layout will not include a layers dimension.
103 +
104 + Returns:
105 + A tuple of ints which is a permutation of range(len(shape)).
106 + """
107 raise NotImplementedError
108
109 @classmethod
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+31 -2
@@ -38,7 +38,7 @@ The class provides the following primitives:
38 import enum
39 from abc import ABC, abstractmethod
40 from collections.abc import Callable, Iterable
41 -from typing import TYPE_CHECKING, Any, Literal, Optional
41 +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional
42
43 import torch
44
@@ -47,7 +47,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
47 from vllm.v1.outputs import KVConnectorOutput
48
49 if TYPE_CHECKING:
50 - from vllm.attention.backends.abstract import AttentionMetadata
50 + from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
51 from vllm.config import VllmConfig
52 from vllm.distributed.kv_events import KVCacheEvent
53 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
@@ -142,6 +142,18 @@ class KVConnectorMetadata(ABC): # noqa: B024
142
143
144 class KVConnectorBase_V1(ABC):
145 + """
146 + Base class for KV connectors.
147 +
148 + Attributes:
149 + prefer_cross_layer_blocks (bool): Indicates whether this connector
150 + prefers KV blocks that hold KV data for all layers (for speeding
151 + up KV data transfers).
152 + Defaults to False.
153 + """
154 +
155 + prefer_cross_layer_blocks: ClassVar[bool] = False
156 +
157 def __init__(
158 self,
159 vllm_config: "VllmConfig",
@@ -226,6 +238,23 @@ class KVConnectorBase_V1(ABC):
238 """
239 return
240
241 + def register_cross_layers_kv_cache(
242 + self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]
243 + ):
244 + """
245 + Initialize with a single KV cache tensor used by all layers.
246 + The first dimension should be num_layers.
247 + This function will only be called for models with uniform layers,
248 + and only if the prefers_cross_layer_blocks is set to True.
249 + Only one of the functions
250 + {register_kv_caches, register_cross_layers_kv_cache} will be called.
251 +
252 + Args:
253 + kv_cache: a cross-layers kv cache tensor
254 + attn_backend: The attention backend that corresponds to all layers
255 + """
256 + return
257 +
258 def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
259 """
260 Set the xPU-specific ops for copying KV between host and device.
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+38 -5
@@ -4,12 +4,12 @@ from collections import defaultdict
4 from collections.abc import Iterable, Iterator
5 from dataclasses import dataclass
6 from itertools import islice
7 -from typing import Any
7 +from typing import Any, ClassVar
8
9 import torch
10
11 -from vllm.attention import AttentionMetadata
12 -from vllm.config import VllmConfig
11 +from vllm.attention import Attention, AttentionBackend, AttentionMetadata
12 +from vllm.config import VllmConfig, get_layers_from_vllm_config
13 from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
14 from vllm.distributed.kv_transfer.kv_connector.v1 import (
15 KVConnectorBase_V1,
@@ -42,6 +42,8 @@ class OffloadingConnectorMetadata(KVConnectorMetadata):
42
43
44 class OffloadingConnector(KVConnectorBase_V1):
45 + prefer_cross_layer_blocks: ClassVar[bool] = True
46 +
47 def __init__(
48 self,
49 vllm_config: VllmConfig,
@@ -63,6 +65,12 @@ class OffloadingConnector(KVConnectorBase_V1):
65 assert self.connector_worker is not None
66 self.connector_worker.register_kv_caches(kv_caches)
67
68 + def register_cross_layers_kv_cache(
69 + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
70 + ):
71 + assert self.connector_worker is not None
72 + self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)
73 +
74 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
75 assert self.connector_worker is not None
76 assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
@@ -422,10 +430,35 @@ class OffloadingConnectorWorker:
430 self._job_counter = job_id + 1
431 return job_id
432
425 - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
426 - for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):
433 + def _register_handlers(
434 + self,
435 + kv_caches: dict[str, torch.Tensor],
436 + attn_backends: dict[str, type[AttentionBackend]],
437 + ):
438 + for src_cls, dst_cls, handler in self.spec.get_handlers(
439 + kv_caches, attn_backends
440 + ):
441 self.worker.register_handler(src_cls, dst_cls, handler)
442
443 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
444 + layer_names = list(kv_caches.keys())
445 + layers = get_layers_from_vllm_config(
446 + self.spec.vllm_config, Attention, layer_names
447 + )
448 + attn_backends = {
449 + layer_name: layers[layer_name].get_attn_backend()
450 + for layer_name in layer_names
451 + }
452 + self._register_handlers(kv_caches, attn_backends)
453 +
454 + def register_cross_layers_kv_cache(
455 + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
456 + ):
457 + cross_layer_name = "ALL_LAYERS"
458 + kv_caches = {cross_layer_name: kv_cache}
459 + attn_backends = {cross_layer_name: attn_backend}
460 + self._register_handlers(kv_caches, attn_backends)
461 +
462 def start_load_kv(self, metadata: OffloadingConnectorMetadata):
463 for req_id, transfer_spec in metadata.reqs_to_load.items():
464 job_id = self._generate_job_id()
vllm/v1/attention/backends/flash_attn.py
+10 -2
@@ -99,12 +99,20 @@ class FlashAttentionBackend(AttentionBackend):
99 return (2, num_blocks, block_size, num_kv_heads, head_size)
100
101 @staticmethod
102 - def get_kv_cache_stride_order() -> tuple[int, ...]:
102 + def get_kv_cache_stride_order(
103 + include_num_layers_dimension: bool = False,
104 + ) -> tuple[int, ...]:
105 # `stride_order` indicates the permutation that gets
106 # us from `get_kv_cache_shape` to the actual memory layout we want.
107 cache_layout = get_kv_cache_layout()
106 - if cache_layout == "NHD":
108 + if cache_layout == "NHD" and include_num_layers_dimension:
109 + # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
110 + return (2, 0, 1, 3, 4, 5)
111 + elif cache_layout == "NHD":
112 stride_order = (0, 1, 2, 3, 4)
113 + elif cache_layout == "HND" and include_num_layers_dimension:
114 + # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
115 + return (2, 4, 0, 1, 3, 5)
116 elif cache_layout == "HND":
117 stride_order = (0, 1, 3, 2, 4)
118 else:
vllm/v1/attention/backends/flashinfer.py
+10 -2
@@ -309,12 +309,20 @@ class FlashInferBackend(AttentionBackend):
309 return (num_blocks, 2, block_size, num_kv_heads, head_size)
310
311 @staticmethod
312 - def get_kv_cache_stride_order() -> tuple[int, ...]:
312 + def get_kv_cache_stride_order(
313 + include_num_layers_dimension: bool = False,
314 + ) -> tuple[int, ...]:
315 # `stride_order` indicates the permutation that gets us from
316 # `get_kv_cache_shape` to the actual memory layout we want.
317 cache_layout = get_kv_cache_layout()
316 - if cache_layout == "NHD":
318 + if cache_layout == "NHD" and include_num_layers_dimension:
319 + # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
320 + return (1, 0, 2, 3, 4, 5)
321 + elif cache_layout == "NHD":
322 stride_order = (0, 1, 2, 3, 4)
323 + elif cache_layout == "HND" and include_num_layers_dimension:
324 + # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
325 + return (1, 2, 4, 0, 3, 5)
326 elif cache_layout == "HND":
327 stride_order = (0, 1, 3, 2, 4)
328 else:
vllm/v1/attention/backends/mla/common.py
+9
@@ -308,6 +308,15 @@ class MLACommonBackend(AttentionBackend):
308 ) -> tuple[int, ...]:
309 return (num_blocks, block_size, head_size)
310
311 + @staticmethod
312 + def get_kv_cache_stride_order(
313 + include_num_layers_dimension: bool = False,
314 + ) -> tuple[int, ...]:
315 + # `stride_order` indicates the permutation that gets
316 + # us from `get_kv_cache_shape` to the actual memory layout we want.
317 + # (num_blocks, num_layers, block_size, head_size)
318 + return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2)
319 +
320 @classmethod
321 def get_supported_head_sizes(cls) -> list[int]:
322 return [576]
vllm/v1/attention/backends/mla/indexer.py
+5 -1
@@ -48,7 +48,11 @@ class DeepseekV32IndexerBackend(AttentionBackend):
48 return (num_blocks, block_size, head_size)
49
50 @staticmethod
51 - def get_kv_cache_stride_order() -> tuple[int, ...]:
51 + def get_kv_cache_stride_order(
52 + include_num_layers_dimension: bool = False,
53 + ) -> tuple[int, ...]:
54 + if include_num_layers_dimension:
55 + return (0, 1, 2, 3)
56 return (0, 1, 2)
57
58
vllm/v1/kv_offload/cpu.py
+5 -12
@@ -4,8 +4,8 @@ from collections.abc import Iterator
4
5 import torch
6
7 -from vllm.config import VllmConfig, get_layers_from_vllm_config
8 -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
7 +from vllm.attention import AttentionBackend
8 +from vllm.config import VllmConfig
9 from vllm.platforms import current_platform
10 from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
11 from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
@@ -63,7 +63,9 @@ class CPUOffloadingSpec(OffloadingSpec):
63 return self._manager
64
65 def get_handlers(
66 - self, kv_caches: dict[str, torch.Tensor]
66 + self,
67 + kv_caches: dict[str, torch.Tensor],
68 + attn_backends: dict[str, type[AttentionBackend]],
69 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
70 if not self._handler:
71 if not current_platform.is_cuda_alike():
@@ -71,15 +73,6 @@ class CPUOffloadingSpec(OffloadingSpec):
73 "CPU Offloading is currently only supported on CUDA-alike GPUs"
74 )
75
74 - layer_names = list(kv_caches.keys())
75 - layers = get_layers_from_vllm_config(
76 - self.vllm_config, AttentionLayerBase, layer_names
77 - )
78 - attn_backends = {
79 - layer_name: layers[layer_name].get_attn_backend()
80 - for layer_name in layer_names
81 - }
82 -
76 self._handler = CpuGpuOffloadingHandler(
77 attn_backends=attn_backends,
78 gpu_block_size=self.gpu_block_size,
vllm/v1/kv_offload/spec.py
+5 -1
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
11 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
12
13 if TYPE_CHECKING:
14 + from vllm.attention import AttentionBackend
15 from vllm.config import VllmConfig
16
17 logger = init_logger(__name__)
@@ -48,13 +49,16 @@ class OffloadingSpec(ABC):
49
50 @abstractmethod
51 def get_handlers(
51 - self, kv_caches: dict[str, torch.Tensor]
52 + self,
53 + kv_caches: dict[str, torch.Tensor],
54 + attn_backends: dict[str, type["AttentionBackend"]],
55 ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
56 """
57 Get offloading handlers along with their respective src and dst types.
58
59 Args:
60 kv_caches: A dictionary of layer_name -> gpu_kv_cache tensor.
61 + attn_backends: A dictionary of layer_name -> AttentionBackend.
62
63 Yields:
64 Tuples of (src_type, dst_type, offloading_handler).
vllm/v1/kv_offload/worker/cpu_gpu.py
+10 -2
@@ -83,10 +83,18 @@ class CpuGpuOffloadingHandler(OffloadingHandler):
83 self.gpu_tensors.append(gpu_tensor)
84
85 gpu_shape = gpu_tensor.shape
86 - test_shape = attn_backends[layer_name].get_kv_cache_shape(
86 + attn_backend = attn_backends[layer_name]
87 + test_shape = attn_backend.get_kv_cache_shape(
88 num_blocks=1234, block_size=16, num_kv_heads=8, head_size=256
89 )
89 - if test_shape[0] == 1234:
90 +
91 + if len(gpu_shape) != len(test_shape):
92 + # cross-layers tensor
93 + # shape is (num_blocks, ...)
94 + assert len(gpu_shape) == len(test_shape) + 1
95 + num_blocks_idx = 0
96 + self.kv_dim_before_num_blocks.append(False)
97 + elif test_shape[0] == 1234:
98 # shape is (num_blocks, ...)
99 num_blocks_idx = 0
100 self.kv_dim_before_num_blocks.append(False)
vllm/v1/worker/gpu_model_runner.py
+34 -7
@@ -349,6 +349,9 @@ class GPUModelRunner(
349 # self.model: nn.Module # Set after load_model
350 # Initialize in initialize_kv_cache
351 self.kv_caches: list[torch.Tensor] = []
352 + # Initialize in initialize_kv_cache_tensors
353 + self.cross_layers_kv_cache: torch.Tensor | None = None
354 + self.cross_layers_attn_backend: type[AttentionBackend] | None = None
355 # indexes: [kv_cache_group_id][attn_group]
356 self.attn_groups: list[list[AttentionGroup]] = []
357 # self.kv_cache_config: KVCacheConfig
@@ -4930,12 +4933,30 @@ class GPUModelRunner(
4933 Dict[str, torch.Tensor]: A map between layer names to their
4934 corresponding memory buffer for KV cache.
4935 """
4933 - # Initialize the memory buffer for KV cache
4934 - kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
4935 - # Change the memory buffer to the desired shape
4936 - kv_caches = self._reshape_kv_cache_tensors(
4937 - kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes
4938 - )
4936 +
4937 + # Try creating KV caches optimized for kv-connector transfers
4938 + cache_dtype = self.cache_config.cache_dtype
4939 + if self.use_uniform_kv_cache(self.attn_groups, cache_dtype):
4940 + kv_caches, cross_layers_kv_cache, attn_backend = (
4941 + self.allocate_uniform_kv_caches(
4942 + kv_cache_config,
4943 + self.attn_groups,
4944 + cache_dtype,
4945 + self.device,
4946 + kernel_block_sizes,
4947 + )
4948 + )
4949 + self.cross_layers_kv_cache = cross_layers_kv_cache
4950 + self.cross_layers_attn_backend = attn_backend
4951 + else:
4952 + # Fallback to the general case
4953 + # Initialize the memory buffer for KV cache
4954 + kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
4955 +
4956 + # Change the memory buffer to the desired shape
4957 + kv_caches = self._reshape_kv_cache_tensors(
4958 + kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes
4959 + )
4960
4961 # Set up cross-layer KV cache sharing
4962 for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
@@ -5017,7 +5038,13 @@ class GPUModelRunner(
5038
5039 if has_kv_transfer_group():
5040 kv_transfer_group = get_kv_transfer_group()
5020 - kv_transfer_group.register_kv_caches(kv_caches)
5041 + if self.cross_layers_kv_cache is not None:
5042 + assert self.cross_layers_attn_backend is not None
5043 + kv_transfer_group.register_cross_layers_kv_cache(
5044 + self.cross_layers_kv_cache, self.cross_layers_attn_backend
5045 + )
5046 + else:
5047 + kv_transfer_group.register_kv_caches(kv_caches)
5048 kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
5049
5050 if self.dcp_world_size > 1:
vllm/v1/worker/kv_connector_model_runner_mixin.py
+165
@@ -11,7 +11,11 @@ from typing import (
11 TYPE_CHECKING, # noqa: UP035
12 )
13
14 +import torch
15 +
16 +from vllm.attention import AttentionBackend
17 from vllm.config import VllmConfig
18 +from vllm.config.cache import CacheDType
19 from vllm.distributed.kv_transfer import (
20 ensure_kv_transfer_shutdown,
21 get_kv_transfer_group,
@@ -21,11 +25,13 @@ from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
25 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
26 from vllm.forward_context import get_forward_context, set_forward_context
27 from vllm.logger import init_logger
28 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
29 from vllm.v1.outputs import (
30 EMPTY_MODEL_RUNNER_OUTPUT,
31 KVConnectorOutput,
32 ModelRunnerOutput,
33 )
34 +from vllm.v1.worker.utils import AttentionGroup
35
36 if TYPE_CHECKING:
37 from vllm.v1.core.sched.output import SchedulerOutput
@@ -142,3 +148,162 @@ class KVConnectorModelRunnerMixin:
148 if has_kv_transfer_group():
149 return get_kv_transfer_group().get_kv_connector_stats()
150 return None
151 +
152 + @staticmethod
153 + def use_uniform_kv_cache(
154 + attn_groups: list[list[AttentionGroup]],
155 + cache_dtype: CacheDType,
156 + ) -> bool:
157 + """
158 + Determines whether a uniform KV layout should be used.
159 + A uniform layout means all layers KV caches will share the same
160 + underlying tensor, where for a given block number, the respective
161 + KV data for all layers will be contiguous.
162 + This will allow efficient KV transfer of per-block KV data for all
163 + layers at once.
164 + Note this layout will only be applied given 3 conditions:
165 + 1. The KV Cache config contains just a single group where all layers
166 + have the same page size.
167 + 2. A KV connector is configured, and the KV connector instance prefers
168 + to use this layout (prefer_cross_layer_blocks() returns True)
169 + 2. The flash attention backend supports this layout
170 + (get_kv_cache_stride_order(True) includes a placement for a
171 + num_layers dimension)
172 +
173 + Note that the actual placement of the num_layers dimensions
174 + in the unified layers tensors will be determined by the attention
175 + backend.
176 + Thus, the layers KV data may still not be contiguous per block
177 + if the attention backend does not support it.
178 +
179 + Args:
180 + attn_groups: The list of attention groups for this model
181 + cache_dtype: The KV cache dtype
182 + Returns:
183 + True if we should use a uniform KV cache layout.
184 + """
185 +
186 + if not has_kv_transfer_group():
187 + return False
188 + if not get_kv_transfer_group().prefer_cross_layer_blocks:
189 + return False
190 +
191 + if len(attn_groups) != 1 or len(attn_groups[0]) != 1:
192 + return False
193 +
194 + attn_group = attn_groups[0][0]
195 + kv_cache_spec = attn_group.kv_cache_spec
196 + if not isinstance(kv_cache_spec, AttentionSpec):
197 + return False
198 +
199 + attn_backend = attn_group.backend
200 + kv_cache_shape = attn_backend.get_kv_cache_shape(
201 + 1234,
202 + kv_cache_spec.block_size,
203 + kv_cache_spec.num_kv_heads,
204 + kv_cache_spec.head_size,
205 + cache_dtype_str=cache_dtype,
206 + )
207 +
208 + try:
209 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
210 + include_num_layers_dimension=True
211 + )
212 + except (AttributeError, NotImplementedError):
213 + return False
214 +
215 + # check that attention backend include a layers dimension
216 + return len(kv_cache_stride_order) == len(kv_cache_shape) + 1
217 +
218 + @staticmethod
219 + def allocate_uniform_kv_caches(
220 + kv_cache_config: KVCacheConfig,
221 + attn_groups: list[list[AttentionGroup]],
222 + cache_dtype: CacheDType,
223 + device: torch.device,
224 + kernel_block_sizes: list[int],
225 + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, type[AttentionBackend]]:
226 + """
227 + Initializes and reshapes KV caches for the simple case where all
228 + layers have the same layout.
229 +
230 + This function assumes use_uniform_kv_cache() returned True.
231 +
232 + Args:
233 + kv_cache_config: The KV cache config
234 + attn_groups: The list of attention groups for this model
235 + cache_dtype: The KV cache dtype
236 + device: The torch device to allocate on.
237 + kernel_block_sizes: The kernel block sizes for each KV cache group.
238 + Returns:
239 + A tuple (kv_caches, cross_layers_kv_cache, attn_backend) where:
240 + kv_caches is a dict mapping between layer names to their
241 + corresponding memory buffer for KV cache.
242 + cross_layers_kv_cache is the cross layers kv cache tensor
243 + attn_backend is the attention backend matching this tensor
244 + """
245 + attn_group = attn_groups[0][0]
246 + kv_cache_spec = attn_group.kv_cache_spec
247 + assert isinstance(kv_cache_spec, AttentionSpec)
248 +
249 + tensor_sizes = set(
250 + kv_cache_tensor.size for kv_cache_tensor in kv_cache_config.kv_cache_tensors
251 + )
252 + assert len(tensor_sizes) == 1
253 + tensor_size = tensor_sizes.pop()
254 +
255 + page_size = kv_cache_spec.page_size_bytes
256 + assert tensor_size % page_size == 0
257 + num_blocks = tensor_size // page_size
258 + num_layers = len(kv_cache_config.kv_cache_tensors)
259 + total_size = tensor_size * num_layers
260 +
261 + assert len(kernel_block_sizes) == 1
262 + kernel_block_size = kernel_block_sizes[0]
263 + num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size
264 + kernel_num_blocks = num_blocks * num_blocks_per_kv_block
265 +
266 + attn_backend = attn_group.backend
267 + kv_cache_shape = attn_backend.get_kv_cache_shape(
268 + kernel_num_blocks,
269 + kernel_block_size,
270 + kv_cache_spec.num_kv_heads,
271 + kv_cache_spec.head_size,
272 + cache_dtype_str=cache_dtype,
273 + )
274 +
275 + # prepend a num_layers dimension into the shape
276 + kv_cache_shape = (num_layers,) + kv_cache_shape
277 +
278 + try:
279 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
280 + include_num_layers_dimension=True
281 + )
282 + assert len(kv_cache_stride_order) == len(kv_cache_shape)
283 + except (AttributeError, NotImplementedError):
284 + kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
285 +
286 + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
287 +
288 + logger.info("Allocating a cross layer KV cache of shape %s", kv_cache_shape)
289 +
290 + # allocate one contiguous buffer for all layers
291 + cross_layers_kv_cache = (
292 + torch.zeros(total_size, dtype=torch.int8, device=device)
293 + .view(kv_cache_spec.dtype)
294 + .view(kv_cache_shape)
295 + )
296 +
297 + # Maintain original KV shape view.
298 + inv_order = [
299 + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
300 + ]
301 + permuted_kv_cache = cross_layers_kv_cache.permute(*inv_order)
302 +
303 + kv_caches = {}
304 + for i, kv_cache_tensor in enumerate(kv_cache_config.kv_cache_tensors):
305 + tensor = permuted_kv_cache[i]
306 + for layer_name in kv_cache_tensor.shared_by:
307 + kv_caches[layer_name] = tensor
308 +
309 + return kv_caches, cross_layers_kv_cache, attn_backend