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 8, 20:15 UTC
Enable Cross layers KV cache layout at NIXL Connector V2 (#33339)
Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
liranschour committed
Feb 5, 2026 at 12:17 UTC
8322d4e47f89f7985b9b3b808fc4ba8549d6afcd
6 files changed
+339
-89
docs/features/nixl_connector_usage.md
+9
@@ -213,6 +213,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
213
--kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
214
```
215
216
+### Cross layers blocks
217
+
218
+By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
219
+To enable this feature:
220
+
221
+```bash
222
+--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
223
+```
224
+
225
## Example Scripts/Code
226
227
Refer to these example scripts in the vLLM repository:
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+8
-2
@@ -14,8 +14,8 @@ tp_configs=(
14
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
15
)
16
dp_ep_configs=(
17
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
18
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
17
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
18
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
19
)
20
21
# Select config array based on DP_EP env var
@@ -57,3 +57,9 @@ if [[ -n "${FLASHINFER:-}" ]]; then
57
else
58
echo "FLASHINFER not set, skipping FLASHINFER runs."
59
fi
60
+
61
+# Check if cross-layers is enabled (non-empty)
62
+if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
63
+ echo "CROSS_LAYERS_BLOCKS is set, rerunning with --enable-cross-layers"
64
+ run_tests "default backend" "--enable-cross-layers"
65
+fi
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+13
-2
@@ -4,6 +4,7 @@ set -xe
4
# Parse command line arguments
5
KV_BUFFER_DEVICE="cuda" # Default to cuda
6
ATTENTION_BACKEND="" # Default to empty (use vllm default)
7
+CROSS_LAYERS_BLOCKS="False"
8
while [[ $# -gt 0 ]]; do
9
case $1 in
10
--kv_buffer_device)
@@ -14,6 +15,10 @@ while [[ $# -gt 0 ]]; do
15
ATTENTION_BACKEND="$2"
16
shift 2
17
;;
18
+ --enable-cross-layers)
19
+ CROSS_LAYERS_BLOCKS="True"
20
+ shift 1
21
+ ;;
22
*)
23
echo "Unknown option $1"
24
echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
@@ -34,11 +39,17 @@ else
39
KV_CONFIG_HETERO_LAYOUT=''
40
fi
41
42
+if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
43
+ KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
44
+else
45
+ KV_EXTRA_CONFIG=''
46
+fi
47
+
48
# Build the kv-transfer-config once
49
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
39
- KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
50
+ KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
51
else
41
- KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
52
+ KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
53
fi
54
55
# Models to run
tests/v1/kv_connector/unit/test_nixl_connector.py
+188
-47
@@ -18,8 +18,12 @@ import ray
18
import torch
19
20
from vllm import LLM
21
-from vllm.config import KVTransferConfig
22
-from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
21
+from vllm.config import KVTransferConfig, set_current_vllm_config
22
+from vllm.distributed.kv_transfer.kv_connector.utils import (
23
+ KVOutputAggregator,
24
+ TpKVTopology,
25
+ get_current_attn_backend,
26
+)
27
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
28
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
29
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
@@ -46,10 +50,14 @@ from vllm.platforms import current_platform
50
from vllm.platforms.interface import Platform
51
from vllm.sampling_params import SamplingParams
52
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
53
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
54
from vllm.v1.engine import EngineCoreRequest
55
from vllm.v1.engine.output_processor import OutputProcessor
56
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
57
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
58
from vllm.v1.request import RequestStatus
59
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
60
+from vllm.v1.worker.utils import AttentionGroup
61
62
from .utils import create_request, create_scheduler, create_vllm_config
63
@@ -366,6 +374,7 @@ def test_kv_transfer_handshake(dist_init):
374
375
# Decode connector will be able to create handshake with the prefill connector.
376
decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
377
+ decode_connector.register_kv_caches(kv_caches)
378
379
# Here we are testing the retrieval of NIXLAgentMetadata.
380
# Knowing the implementation detail, we override the add_remote_agent
@@ -402,6 +411,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
411
self.kv_cache_layout = kv_cache_layout
412
# Mock register_kv_caches attribute needed for tests that do not call it.
413
self.src_xfer_handles_by_block_size = {self.block_size: 1}
414
+ test_shape = self.attn_backend.get_kv_cache_shape(
415
+ num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
416
+ )
417
+ self.kv_topo = TpKVTopology(
418
+ tp_rank=self.tp_rank,
419
+ engine_id=self.engine_id,
420
+ remote_tp_size=self._tp_size, # shared state
421
+ remote_block_size=self._block_size, # shared state
422
+ is_mla=self.use_mla,
423
+ total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
424
+ attn_backend=self.attn_backend,
425
+ tensor_shape=test_shape,
426
+ )
427
+
428
+ self.compat_hash = compute_nixl_compatibility_hash(
429
+ self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
430
+ )
431
432
def _nixl_handshake(
433
self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
@@ -1352,6 +1378,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1378
llm.llm_engine.engine_core.shutdown()
1379
1380
1381
+@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
1382
@pytest.mark.parametrize(
1383
"attn_backend",
1384
[
@@ -1372,7 +1399,9 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1399
"TRITON_ATTN",
1400
],
1401
)
1375
-def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1402
+def test_register_kv_caches(
1403
+ default_vllm_config, dist_init, attn_backend, enable_cross_layers
1404
+):
1405
"""
1406
Test that register_kv_caches() properly calls nixl_wrapper methods with
1407
correct data.
@@ -1386,6 +1415,12 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1415
1416
vllm_config = create_vllm_config(attention_backend=attn_backend)
1417
1418
+ # Enable cross layers blocks
1419
+ vllm_config.kv_transfer_config.kv_connector_extra_config[
1420
+ "enable_cross_layers_blocks"
1421
+ ] = enable_cross_layers
1422
+ set_kv_cache_layout("HND")
1423
+
1424
# Import the appropriate backend based on the parameter
1425
if attn_backend == "FLASH_ATTN":
1426
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
@@ -1400,44 +1435,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1435
1436
backend_cls = TritonAttentionBackend
1437
1403
- # Create test kv cache tensors using proper backend shape
1404
- kv_cache_shape = backend_cls.get_kv_cache_shape(
1405
- num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1406
- )
1407
- shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1408
- unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1409
- kv_caches = {
1410
- "layer0": shared_tensor,
1411
- "layer1": unique_tensor,
1412
- "layer2": shared_tensor,
1413
- }
1414
-
1415
- # Store tensor info for validation
1416
-
1417
- test_shape = backend_cls.get_kv_cache_shape(
1418
- num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1419
- )
1420
- is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1421
-
1422
- if is_blocks_first:
1423
- expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
1424
- expected_base_addrs = [
1425
- shared_tensor.data_ptr(),
1426
- unique_tensor.data_ptr(),
1427
- ]
1428
- expected_num_entries = 2
1429
- else:
1430
- expected_tensor_size = (
1431
- shared_tensor[0].element_size() * shared_tensor[0].numel()
1432
- )
1433
- expected_base_addrs = [
1434
- shared_tensor[0].data_ptr(),
1435
- shared_tensor[1].data_ptr(),
1436
- unique_tensor[0].data_ptr(),
1437
- unique_tensor[1].data_ptr(),
1438
- ]
1439
- expected_num_entries = 4
1440
-
1438
nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1439
with (
1440
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
@@ -1466,6 +1463,111 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1463
# Reassure the shutdown() check that the thread is terminated
1464
mock_thread.return_value.is_alive.return_value = False
1465
1466
+ expected_tensor_size: int
1467
+ expected_base_addrs: list[int]
1468
+ expected_num_entries: int
1469
+ kv_caches: dict[str, torch.Tensor]
1470
+ assert str(enable_cross_layers).lower() != "true" or (
1471
+ (attn_backend not in ("FLASH_ATTN", "FLASHINFER"))
1472
+ or connector.prefer_cross_layer_blocks
1473
+ )
1474
+ if connector.prefer_cross_layer_blocks:
1475
+ num_layers = 32
1476
+ block_size = 16
1477
+ num_blocks = 8
1478
+ kv_cache_spec = AttentionSpec(
1479
+ block_size=block_size,
1480
+ num_kv_heads=4,
1481
+ head_size=64,
1482
+ dtype=torch.bfloat16,
1483
+ )
1484
+ kv_cache_config = KVCacheConfig(
1485
+ num_blocks=num_blocks,
1486
+ kv_cache_tensors=[
1487
+ KVCacheTensor(
1488
+ size=kv_cache_spec.page_size_bytes * num_blocks,
1489
+ shared_by=["dummy-layer"],
1490
+ )
1491
+ for i in range(num_layers)
1492
+ ],
1493
+ # allocate_uniform_kv_caches does not use this
1494
+ kv_cache_groups=[],
1495
+ )
1496
+
1497
+ with set_current_vllm_config(vllm_config):
1498
+ _, cross_layers_kv_cache, _ = (
1499
+ KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
1500
+ kv_cache_config=kv_cache_config,
1501
+ attn_groups=[
1502
+ [
1503
+ AttentionGroup(
1504
+ backend=backend_cls,
1505
+ layer_names=[],
1506
+ kv_cache_spec=kv_cache_spec,
1507
+ kv_cache_group_id=0,
1508
+ )
1509
+ ]
1510
+ ],
1511
+ cache_dtype=torch.bfloat16,
1512
+ device=torch.cuda.current_device(),
1513
+ kernel_block_sizes=[block_size],
1514
+ )
1515
+ )
1516
+ # Store tensor info for validation
1517
+ expected_tensor_size = (
1518
+ cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
1519
+ )
1520
+ expected_base_addrs = [
1521
+ cross_layers_kv_cache.data_ptr(),
1522
+ ]
1523
+ expected_num_entries = 1
1524
+
1525
+ expected_blocks_count = 8
1526
+
1527
+ kv_caches = {"all-layers": cross_layers_kv_cache}
1528
+
1529
+ else:
1530
+ # Create test kv cache tensors using proper backend shape
1531
+ kv_cache_shape = backend_cls.get_kv_cache_shape(
1532
+ num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1533
+ )
1534
+ shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1535
+ unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1536
+ kv_caches = {
1537
+ "layer0": shared_tensor,
1538
+ "layer1": unique_tensor,
1539
+ "layer2": shared_tensor,
1540
+ }
1541
+
1542
+ # Store tensor info for validation
1543
+
1544
+ test_shape = backend_cls.get_kv_cache_shape(
1545
+ num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1546
+ )
1547
+ is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1548
+
1549
+ if is_blocks_first:
1550
+ expected_tensor_size = (
1551
+ shared_tensor.element_size() * shared_tensor.numel()
1552
+ )
1553
+ expected_base_addrs = [
1554
+ shared_tensor.data_ptr(),
1555
+ unique_tensor.data_ptr(),
1556
+ ]
1557
+ expected_num_entries = 2
1558
+ else:
1559
+ expected_tensor_size = (
1560
+ shared_tensor[0].element_size() * shared_tensor[0].numel()
1561
+ )
1562
+ expected_base_addrs = [
1563
+ shared_tensor[0].data_ptr(),
1564
+ shared_tensor[1].data_ptr(),
1565
+ unique_tensor[0].data_ptr(),
1566
+ unique_tensor[1].data_ptr(),
1567
+ ]
1568
+ expected_num_entries = 4
1569
+ expected_blocks_count = 8
1570
+
1571
# Execute register_kv_caches
1572
connector.register_kv_caches(kv_caches)
1573
@@ -1489,16 +1591,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1591
blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
1592
1593
# Validate blocks_data structure and size
1492
- expected_blocks_count = 8
1594
assert len(blocks_data) == expected_blocks_count, (
1595
f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
1596
)
1597
1497
- num_blocks = 2
1498
- if is_blocks_first:
1499
- expected_block_len = expected_tensor_size // num_blocks // 2
1500
- else:
1598
+ if connector.prefer_cross_layer_blocks:
1599
+ num_blocks = 8
1600
expected_block_len = expected_tensor_size // num_blocks
1601
+ else:
1602
+ num_blocks = 2
1603
+ if is_blocks_first:
1604
+ expected_block_len = expected_tensor_size // num_blocks // 2
1605
+ else:
1606
+ expected_block_len = expected_tensor_size // num_blocks
1607
1608
for i, block_entry in enumerate(blocks_data):
1609
block_start_addr, block_len, tp_rank = block_entry
@@ -1507,6 +1612,8 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1612
f"got {block_len}"
1613
)
1614
1615
+ assert connector.connector_worker.block_size == 16
1616
+
1617
1618
class FakePlatform(Platform):
1619
device_type: str = "oot"
@@ -2049,6 +2156,17 @@ def test_compatibility_hash_validation(
2156
)
2157
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2158
decode_worker = decode_connector.connector_worker
2159
+ kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
2160
+ num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
2161
+ )
2162
+ shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2163
+ unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2164
+ kv_caches = {
2165
+ "layer0": shared_tensor,
2166
+ "layer1": unique_tensor,
2167
+ "layer2": shared_tensor,
2168
+ }
2169
+ decode_connector.register_kv_caches(kv_caches)
2170
2171
remote_config_params: dict[str, Any] = {
2172
"model": "facebook/opt-125m",
@@ -2071,7 +2189,9 @@ def test_compatibility_hash_validation(
2189
)
2190
)
2191
remote_hash = compute_nixl_compatibility_hash(
2074
- remote_vllm_config, decode_worker.backend_name
2192
+ remote_vllm_config,
2193
+ decode_worker.backend_name,
2194
+ decode_worker.kv_topo.cross_layers_blocks,
2195
)
2196
2197
prefill_block_size = config_overrides.get("block_size", 16)
@@ -2150,6 +2270,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2270
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2271
decode_worker = decode_connector.connector_worker
2272
2273
+ backend = get_current_attn_backend(local_vllm_config)
2274
+ test_shape = backend.get_kv_cache_shape(
2275
+ num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
2276
+ )
2277
+ decode_worker.kv_topo = TpKVTopology(
2278
+ tp_rank=decode_worker.tp_rank,
2279
+ engine_id=decode_worker.engine_id,
2280
+ remote_tp_size=decode_worker._tp_size, # shared state
2281
+ remote_block_size=decode_worker._block_size, # shared state
2282
+ is_mla=decode_worker.use_mla,
2283
+ total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
2284
+ attn_backend=backend,
2285
+ tensor_shape=test_shape,
2286
+ )
2287
+
2288
+ decode_worker.compat_hash = compute_nixl_compatibility_hash(
2289
+ decode_worker.vllm_config,
2290
+ decode_worker.backend_name,
2291
+ decode_worker.kv_topo.cross_layers_blocks,
2292
+ )
2293
+
2294
if error_scenario == "handshake_decode_error":
2295
msg_bytes = b"this is not valid msgpack data"
2296
elif error_scenario == "handshake_validation_error":
vllm/distributed/kv_transfer/kv_connector/utils.py
+45
-4
@@ -14,6 +14,7 @@ from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vll
14
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17
+from vllm.platforms import current_platform
18
from vllm.v1.attention.backend import AttentionBackend
19
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
20
@@ -192,8 +193,6 @@ def copy_kv_blocks(
193
dst_device=dst_device,
194
)
195
195
- from vllm.platforms import current_platform
196
-
196
if direction == "h2d":
197
copy_fn = current_platform.insert_blocks_to_device
198
else:
@@ -316,12 +315,14 @@ class TpKVTopology:
315
attn_backend: type[AttentionBackend]
316
engine_id: EngineId
317
remote_block_size: dict[EngineId, int]
318
+ tensor_shape: torch.Size | None = None
319
320
def __post_init__(self):
321
# Figure out whether the first dimension of the cache is K/V
322
# or num_blocks. This is used to register the memory regions correctly.
323
+ _MOCK_BLOCK_SIZE = 16
324
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
324
- num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
325
+ num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1
326
)
327
# Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
328
# we just mock num_blocks to 1 for the dimension check below.
@@ -329,6 +330,36 @@ class TpKVTopology:
330
len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
331
)
332
333
+ self._cross_layers_blocks = False
334
+ if self.tensor_shape is not None:
335
+ self._cross_layers_blocks = (
336
+ len(self.tensor_shape) == len(kv_cache_shape) + 1
337
+ )
338
+
339
+ if self._cross_layers_blocks:
340
+ # prepend layers dimension
341
+ _MOCK_NUM_LAYERS = 80
342
+ kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
343
+ try:
344
+ kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(
345
+ include_num_layers_dimension=self._cross_layers_blocks
346
+ )
347
+ except (AttributeError, NotImplementedError):
348
+ kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
349
+
350
+ # In case of cross layers permute kv_cache_shape according to
351
+ # stride_order to retrieve physical position of block_size
352
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
353
+
354
+ # In the default non-cross layers layout the block_size position
355
+ # is logical while in the cross layers case it is the physical
356
+ # position. This matches the shape of the actual kv cache tensors
357
+ # passed at register_kv_caches()/register_cross_layers_kv_cache()
358
+ block_size_position = kv_cache_shape.index(_MOCK_BLOCK_SIZE)
359
+
360
+ assert block_size_position is not None
361
+ self._block_size_position = -(len(kv_cache_shape) - block_size_position)
362
+
363
@property
364
def is_kv_layout_blocks_first(self) -> bool:
365
return self._is_kv_layout_blocks_first
@@ -336,7 +367,9 @@ class TpKVTopology:
367
@property
368
def split_k_and_v(self) -> bool:
369
# Whether to register regions for K and V separately (when present).
339
- return not (self.is_mla or self.is_kv_layout_blocks_first)
370
+ return not (
371
+ self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
372
+ )
373
374
@property
375
def tp_size(self) -> int:
@@ -346,6 +379,14 @@ class TpKVTopology:
379
def block_size(self) -> int:
380
return self.remote_block_size[self.engine_id]
381
382
+ @property
383
+ def cross_layers_blocks(self) -> bool:
384
+ return self._cross_layers_blocks
385
+
386
+ @property
387
+ def block_size_position(self) -> int:
388
+ return self._block_size_position
389
+
390
def tp_ratio(
391
self,
392
remote_tp_size: int,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+76
-34
@@ -54,7 +54,7 @@ from vllm.forward_context import ForwardContext
54
from vllm.logger import init_logger
55
from vllm.platforms import current_platform
56
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
57
-from vllm.v1.attention.backend import AttentionMetadata
57
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
58
from vllm.v1.attention.backends.utils import get_kv_cache_layout
59
from vllm.v1.core.sched.output import SchedulerOutput
60
from vllm.v1.worker.block_table import BlockTable
@@ -173,7 +173,7 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata):
173
174
175
def compute_nixl_compatibility_hash(
176
- vllm_config: VllmConfig, attn_backend_name: str
176
+ vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
177
) -> str:
178
"""
179
Compute compatibility hash for NIXL KV transfer.
@@ -216,6 +216,7 @@ def compute_nixl_compatibility_hash(
216
# Attention backend and KV cache dtype affect memory layout
217
"attn_backend_name": attn_backend_name,
218
"cache_dtype": str(cache_config.cache_dtype),
219
+ "cross_layers_blocks": cross_layers_blocks,
220
}
221
222
compat_hash = hash_factors(factors)
@@ -298,6 +299,26 @@ class NixlConnectorMetadata(KVConnectorMetadata):
299
300
301
class NixlConnector(KVConnectorBase_V1):
302
+ @property
303
+ def prefer_cross_layer_blocks(self) -> bool:
304
+ backend = get_current_attn_backend(self._vllm_config)
305
+ if backend.get_name() not in (
306
+ "FLASH_ATTN",
307
+ "FLASHINFER",
308
+ ):
309
+ return False
310
+
311
+ # For now there is no benefit to run cross layers when backend
312
+ # does not support on HND
313
+ if get_kv_cache_layout() != "HND":
314
+ return False
315
+
316
+ extra_config = self.kv_transfer_config.kv_connector_extra_config
317
+ return (
318
+ str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
319
+ == "true"
320
+ )
321
+
322
def __init__(
323
self,
324
vllm_config: VllmConfig,
@@ -309,7 +330,7 @@ class NixlConnector(KVConnectorBase_V1):
330
assert vllm_config.kv_transfer_config is not None
331
assert vllm_config.kv_transfer_config.engine_id is not None
332
self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
312
-
333
+ self.kv_transfer_config = vllm_config.kv_transfer_config
334
if role == KVConnectorRole.SCHEDULER:
335
self.connector_scheduler: NixlConnectorScheduler | None = (
336
NixlConnectorScheduler(vllm_config, self.engine_id)
@@ -395,6 +416,16 @@ class NixlConnector(KVConnectorBase_V1):
416
assert self.connector_worker is not None
417
self.connector_worker.register_kv_caches(kv_caches)
418
419
+ def register_cross_layers_kv_cache(
420
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
421
+ ):
422
+ assert self.connector_worker is not None
423
+
424
+ cross_layer_name = "ALL_LAYERS"
425
+ kv_caches = {cross_layer_name: kv_cache}
426
+
427
+ self.connector_worker.register_kv_caches(kv_caches)
428
+
429
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
430
assert self.connector_worker is not None
431
self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
@@ -976,20 +1007,17 @@ class NixlConnectorWorker:
1007
1008
# Get the attention backend from the first layer
1009
# NOTE (NickLucche) models with multiple backends are not supported yet
979
- backend = get_current_attn_backend(vllm_config)
1010
+ self.attn_backend = get_current_attn_backend(vllm_config)
1011
981
- self.backend_name = backend.get_name()
1012
+ self.backend_name = self.attn_backend.get_name()
1013
self.kv_cache_layout = get_kv_cache_layout()
1014
self.host_buffer_kv_cache_layout = self.kv_cache_layout
1015
logger.debug("Detected attention backend %s", self.backend_name)
1016
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
1017
987
- self.compat_hash = compute_nixl_compatibility_hash(
988
- self.vllm_config, self.backend_name
989
- )
990
- self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
991
- "enforce_handshake_compat", True
992
- )
1018
+ # lazy initialized in register_kv_caches
1019
+ self.compat_hash: str | None = None
1020
+ self.kv_topo: TpKVTopology | None = None
1021
1022
self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
1023
self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
@@ -998,17 +1026,12 @@ class NixlConnectorWorker:
1026
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
1027
self.xfer_stats = NixlKVConnectorStats()
1028
1001
- self.kv_topo = TpKVTopology(
1002
- tp_rank=self.tp_rank,
1003
- engine_id=self.engine_id,
1004
- remote_tp_size=self._tp_size, # shared state
1005
- remote_block_size=self._block_size, # shared state
1006
- is_mla=self.use_mla,
1007
- total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1008
- attn_backend=backend,
1009
- )
1029
self._physical_blocks_per_logical_kv_block = 1
1030
1031
+ self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
1032
+ "enforce_handshake_compat", True
1033
+ )
1034
+
1035
def _nixl_handshake(
1036
self,
1037
host: str,
@@ -1022,6 +1045,7 @@ class NixlConnectorWorker:
1045
# Regardless, only handshake with the remote TP rank(s) that current
1046
# local rank will read from. Note that With homogeneous TP,
1047
# this happens to be the same single rank_i.
1048
+ assert self.kv_topo is not None
1049
p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
1050
remote_rank_to_agent_name = {}
1051
path = make_zmq_path("tcp", host, port)
@@ -1059,6 +1083,7 @@ class NixlConnectorWorker:
1083
)
1084
1085
# Check compatibility hash BEFORE decoding agent metadata
1086
+ assert self.compat_hash is not None
1087
if (
1088
self.enforce_compat_hash
1089
and handshake_payload.compatibility_hash != self.compat_hash
@@ -1267,6 +1292,20 @@ class NixlConnectorWorker:
1292
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1293
"""Register the KV Cache data in nixl."""
1294
1295
+ self.kv_topo = TpKVTopology(
1296
+ tp_rank=self.tp_rank,
1297
+ engine_id=self.engine_id,
1298
+ remote_tp_size=self._tp_size, # shared state
1299
+ remote_block_size=self._block_size, # shared state
1300
+ is_mla=self.use_mla,
1301
+ total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1302
+ attn_backend=self.attn_backend,
1303
+ tensor_shape=next(iter(kv_caches.values())).shape,
1304
+ )
1305
+ self.compat_hash = compute_nixl_compatibility_hash(
1306
+ self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
1307
+ )
1308
+
1309
if self.use_host_buffer:
1310
self.initialize_host_xfer_buffer(kv_caches=kv_caches)
1311
assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -1301,29 +1340,21 @@ class NixlConnectorWorker:
1340
# (roughly 8KB vs 5KB).
1341
# Conversely for FlashInfer, K and V are registered in the same region
1342
# to better exploit the memory layout (ie num_blocks is the first dim).
1304
- split_k_and_v = self.kv_topo.split_k_and_v
1343
tensor_size_bytes = None
1344
1307
- # TODO (NickLucche): Get kernel_block_size in a cleaner way
1308
- # NHD default "view" for non-MLA cache
1309
- if self.device_type == "cpu":
1310
- block_size_position = -2
1311
- else:
1312
- block_size_position = -2 if self.use_mla else -3
1313
-
1345
# Enable different block lengths for different layers when MLA is used.
1346
self.block_len_per_layer = list[int]()
1347
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
1348
for layer_name, cache_or_caches in xfer_buffers.items():
1318
- cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1319
-
1349
+ cache_list = (
1350
+ cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
1351
+ )
1352
for cache in cache_list:
1353
base_addr = cache.data_ptr()
1354
if base_addr in seen_base_addresses:
1355
continue
1356
1325
- kernel_block_size = cache.shape[block_size_position]
1326
-
1357
+ kernel_block_size = cache.shape[self.kv_topo.block_size_position]
1358
if self.block_size != kernel_block_size:
1359
logger.info_once(
1360
"User-specified logical block size (%s) does not match"
@@ -1385,6 +1416,7 @@ class NixlConnectorWorker:
1416
1417
self.device_kv_caches = kv_caches
1418
self.dst_num_blocks[self.engine_id] = self.num_blocks
1419
+
1420
if self.kv_topo.is_kv_layout_blocks_first:
1421
for i in range(len(self.slot_size_per_layer)):
1422
assert self.slot_size_per_layer[i] % 2 == 0
@@ -1440,6 +1472,7 @@ class NixlConnectorWorker:
1472
block_size=self.block_size,
1473
)
1474
# Wrap metadata in payload with hash for defensive decoding
1475
+ assert self.compat_hash is not None
1476
encoder = msgspec.msgpack.Encoder()
1477
self.xfer_handshake_metadata = NixlHandshakePayload(
1478
compatibility_hash=self.compat_hash,
@@ -1461,6 +1494,8 @@ class NixlConnectorWorker:
1494
register another local_xfer_handler using remote block len to ensure
1495
data copy correctness.
1496
"""
1497
+ assert self.kv_topo is not None
1498
+
1499
block_size_ratio = self.block_size // block_size
1500
blocks_data = []
1501
for i, base_addr in enumerate(self.seen_base_addresses):
@@ -1573,6 +1608,7 @@ class NixlConnectorWorker:
1608
# remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
1609
# local origin:| 0| 1| 8| 12|
1610
# local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
1611
+ assert self.kv_topo is not None
1612
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
1613
1614
if engine_id not in self.dst_num_blocks:
@@ -1701,6 +1737,7 @@ class NixlConnectorWorker:
1737
remote_engine_id = nixl_agent_meta.engine_id
1738
1739
assert self._tp_size[remote_engine_id] == remote_tp_size
1740
+ assert self.kv_topo is not None
1741
1742
tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
1743
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
@@ -1837,6 +1874,7 @@ class NixlConnectorWorker:
1874
if len(self.device_kv_caches) == 0:
1875
return
1876
assert block_size_ratio >= 1, "Only nP < nD supported currently."
1877
+ assert self.kv_topo is not None
1878
if self.enable_permute_local_kv and block_size_ratio > 1:
1879
logger.debug(
1880
"Post-processing device kv cache on receive by converting "
@@ -1856,7 +1894,7 @@ class NixlConnectorWorker:
1894
block_size_ratio,
1895
)
1896
1859
- split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
1897
+ split_k_and_v = self.kv_topo.split_k_and_v
1898
1899
for block_ids in block_ids_list:
1900
indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
@@ -1881,6 +1919,7 @@ class NixlConnectorWorker:
1919
The scheduler process (via the MultiprocExecutor) will use this output
1920
to track which workers are done.
1921
"""
1922
+ assert self.kv_topo is not None
1923
done_sending = self._get_new_notifs()
1924
done_recving = self._pop_done_transfers(self._recving_transfers)
1925
@@ -1950,6 +1989,7 @@ class NixlConnectorWorker:
1989
are reading from the same producer (heterogeneous TP scenario), wait
1990
for all consumers to be done pulling.
1991
"""
1992
+ assert self.kv_topo is not None
1993
notified_req_ids: set[str] = set()
1994
for notifs in self.nixl_wrapper.get_new_notifs().values():
1995
for notif in notifs:
@@ -2109,7 +2149,7 @@ class NixlConnectorWorker:
2149
self._reqs_to_send[req_id] = expiration_time
2150
2151
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2112
- assert meta.remote is not None
2152
+ assert meta.remote is not None and self.kv_topo is not None
2153
remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
2154
meta.remote.engine_id
2155
)
@@ -2182,6 +2222,7 @@ class NixlConnectorWorker:
2222
Post a READ point-to-point xfer request from a single local worker to
2223
a single remote worker.
2224
"""
2225
+ assert self.kv_topo is not None
2226
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
2227
if block_size_ratio > 1:
2228
local_block_ids = self.get_mapped_blocks(
@@ -2414,6 +2455,7 @@ class NixlConnectorWorker:
2455
For FlashInfer, this is half the length of the whole block, as K and V
2456
share the same region.
2457
"""
2458
+ assert self.kv_topo is not None
2459
if self.kv_topo.is_kv_layout_blocks_first:
2460
# For indexing only half (either just the K or V part).
2461
block_len = self.block_len_per_layer[layer_idx] // 2