@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 64e3d67ac0
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 9, 19:30 UTC

Enable Cross layers KV cache layout at NIXL Connector (#30207)

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>

liranschour committed Jan 22, 2026 at 12:12 UTC 64e3d67ac01a279b67ef1b67df16882acd74f8fb
5 files changed +308 -89
docs/features/nixl_connector_usage.md
+9
@@ -184,6 +184,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
184 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
185 ```
186
187 +### Cross layers blocks
188 +
189 +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.
190 +To enable this feature:
191 +
192 +```bash
193 +--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
194 +```
195 +
196 ## Example Scripts/Code
197
198 Refer to these example scripts in the vLLM repository:
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+9 -2
@@ -34,11 +34,18 @@ else
34 KV_CONFIG_HETERO_LAYOUT=''
35 fi
36
37 +CROSS_LAYERS_BLOCKS=${CROSS_LAYERS_BLOCKS:-"False"} # Default to non cross layers
38 +if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
39 + KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"cross_layers_blocks": "True"}'
40 +else
41 + KV_EXTRA_CONFIG=''
42 +fi
43 +
44 # Build the kv-transfer-config once
45 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
39 - KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
46 + KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
47 else
41 - KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
48 + KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
49 fi
50
51 # Models to run
tests/v1/kv_connector/unit/test_nixl_connector.py
+178 -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 (
@@ -48,8 +52,11 @@ from vllm.sampling_params import SamplingParams
52 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
53 from vllm.v1.engine import EngineCoreRequest
54 from vllm.v1.engine.output_processor import OutputProcessor
55 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
56 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
57 from vllm.v1.request import RequestStatus
58 +from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
59 +from vllm.v1.worker.utils import AttentionGroup
60
61 from .utils import create_request, create_scheduler, create_vllm_config
62
@@ -366,6 +373,7 @@ def test_kv_transfer_handshake(dist_init):
373
374 # Decode connector will be able to create handshake with the prefill connector.
375 decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
376 + decode_connector.register_kv_caches(kv_caches)
377
378 # Here we are testing the retrieval of NIXLAgentMetadata.
379 # Knowing the implementation detail, we override the add_remote_agent
@@ -402,6 +410,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
410 self.kv_cache_layout = kv_cache_layout
411 # Mock register_kv_caches attribute needed for tests that do not call it.
412 self.src_xfer_handles_by_block_size = {self.block_size: 1}
413 + test_shape = self.attn_backend.get_kv_cache_shape(
414 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
415 + )
416 + self.kv_topo = TpKVTopology(
417 + tp_rank=self.tp_rank,
418 + engine_id=self.engine_id,
419 + remote_tp_size=self._tp_size, # shared state
420 + remote_block_size=self._block_size, # shared state
421 + is_mla=self.use_mla,
422 + total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
423 + attn_backend=self.attn_backend,
424 + tensor_shape=test_shape,
425 + )
426 +
427 + self.compat_hash = compute_nixl_compatibility_hash(
428 + self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
429 + )
430
431 def _nixl_handshake(
432 self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
@@ -1366,6 +1391,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1391 ),
1392 ),
1393 "TRITON_ATTN",
1394 + "FLASHINFER",
1395 ],
1396 )
1397 def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
@@ -1382,6 +1408,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1408
1409 vllm_config = create_vllm_config(attention_backend=attn_backend)
1410
1411 + # Enable cross layers blocks
1412 + vllm_config.kv_transfer_config.kv_connector_extra_config[
1413 + "enable_cross_layers_blocks"
1414 + ] = True
1415 +
1416 # Import the appropriate backend based on the parameter
1417 if attn_backend == "FLASH_ATTN":
1418 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
@@ -1391,49 +1422,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1422 from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
1423
1424 backend_cls = RocmAttentionBackend
1394 - else: # TRITON_ATTN
1425 + else: # TRITON
1426 from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
1427
1428 backend_cls = TritonAttentionBackend
1429
1399 - # Create test kv cache tensors using proper backend shape
1400 - kv_cache_shape = backend_cls.get_kv_cache_shape(
1401 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1402 - )
1403 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1404 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1405 - kv_caches = {
1406 - "layer0": shared_tensor,
1407 - "layer1": unique_tensor,
1408 - "layer2": shared_tensor,
1409 - }
1410 -
1411 - # Store tensor info for validation
1412 -
1413 - test_shape = backend_cls.get_kv_cache_shape(
1414 - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1415 - )
1416 - is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1417 -
1418 - if is_blocks_first:
1419 - expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
1420 - expected_base_addrs = [
1421 - shared_tensor.data_ptr(),
1422 - unique_tensor.data_ptr(),
1423 - ]
1424 - expected_num_entries = 2
1425 - else:
1426 - expected_tensor_size = (
1427 - shared_tensor[0].element_size() * shared_tensor[0].numel()
1428 - )
1429 - expected_base_addrs = [
1430 - shared_tensor[0].data_ptr(),
1431 - shared_tensor[1].data_ptr(),
1432 - unique_tensor[0].data_ptr(),
1433 - unique_tensor[1].data_ptr(),
1434 - ]
1435 - expected_num_entries = 4
1436 -
1430 nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1431 with (
1432 patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
@@ -1462,6 +1455,107 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1455 # Reassure the shutdown() check that the thread is terminated
1456 mock_thread.return_value.is_alive.return_value = False
1457
1458 + expected_tensor_size: int
1459 + expected_base_addrs: list[int]
1460 + expected_num_entries: int
1461 + kv_caches: dict[str, torch.Tensor]
1462 + if connector.prefer_cross_layer_blocks:
1463 + num_layers = 32
1464 + block_size = 16
1465 + num_blocks = 8
1466 + kv_cache_spec = AttentionSpec(
1467 + block_size=block_size,
1468 + num_kv_heads=4,
1469 + head_size=64,
1470 + dtype=torch.bfloat16,
1471 + )
1472 + kv_cache_config = KVCacheConfig(
1473 + num_blocks=num_blocks,
1474 + kv_cache_tensors=[
1475 + KVCacheTensor(
1476 + size=kv_cache_spec.page_size_bytes * num_blocks,
1477 + shared_by=["dummy-layer"],
1478 + )
1479 + for i in range(num_layers)
1480 + ],
1481 + # allocate_uniform_kv_caches does not use this
1482 + kv_cache_groups=[],
1483 + )
1484 +
1485 + with set_current_vllm_config(vllm_config):
1486 + _, cross_layers_kv_cache, _ = (
1487 + KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
1488 + kv_cache_config=kv_cache_config,
1489 + attn_groups=[
1490 + [
1491 + AttentionGroup(
1492 + backend=backend_cls,
1493 + layer_names=[],
1494 + kv_cache_spec=kv_cache_spec,
1495 + kv_cache_group_id=0,
1496 + )
1497 + ]
1498 + ],
1499 + cache_dtype=torch.bfloat16,
1500 + device=torch.cuda.current_device(),
1501 + kernel_block_sizes=[block_size],
1502 + )
1503 + )
1504 + # Store tensor info for validation
1505 + expected_tensor_size = (
1506 + cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
1507 + )
1508 + expected_base_addrs = [
1509 + cross_layers_kv_cache.data_ptr(),
1510 + ]
1511 + expected_num_entries = 1
1512 +
1513 + expected_blocks_count = 8
1514 +
1515 + kv_caches = {"all-layers": cross_layers_kv_cache}
1516 +
1517 + else:
1518 + # Create test kv cache tensors using proper backend shape
1519 + kv_cache_shape = backend_cls.get_kv_cache_shape(
1520 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1521 + )
1522 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1523 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1524 + kv_caches = {
1525 + "layer0": shared_tensor,
1526 + "layer1": unique_tensor,
1527 + "layer2": shared_tensor,
1528 + }
1529 +
1530 + # Store tensor info for validation
1531 +
1532 + test_shape = backend_cls.get_kv_cache_shape(
1533 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1534 + )
1535 + is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1536 +
1537 + if is_blocks_first:
1538 + expected_tensor_size = (
1539 + shared_tensor.element_size() * shared_tensor.numel()
1540 + )
1541 + expected_base_addrs = [
1542 + shared_tensor.data_ptr(),
1543 + unique_tensor.data_ptr(),
1544 + ]
1545 + expected_num_entries = 2
1546 + else:
1547 + expected_tensor_size = (
1548 + shared_tensor[0].element_size() * shared_tensor[0].numel()
1549 + )
1550 + expected_base_addrs = [
1551 + shared_tensor[0].data_ptr(),
1552 + shared_tensor[1].data_ptr(),
1553 + unique_tensor[0].data_ptr(),
1554 + unique_tensor[1].data_ptr(),
1555 + ]
1556 + expected_num_entries = 4
1557 + expected_blocks_count = 8
1558 +
1559 # Execute register_kv_caches
1560 connector.register_kv_caches(kv_caches)
1561
@@ -1485,16 +1579,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1579 blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
1580
1581 # Validate blocks_data structure and size
1488 - expected_blocks_count = 8
1582 assert len(blocks_data) == expected_blocks_count, (
1583 f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
1584 )
1585
1493 - num_blocks = 2
1494 - if is_blocks_first:
1495 - expected_block_len = expected_tensor_size // num_blocks // 2
1496 - else:
1586 + if connector.prefer_cross_layer_blocks:
1587 + num_blocks = 8
1588 expected_block_len = expected_tensor_size // num_blocks
1589 + else:
1590 + num_blocks = 2
1591 + if is_blocks_first:
1592 + expected_block_len = expected_tensor_size // num_blocks // 2
1593 + else:
1594 + expected_block_len = expected_tensor_size // num_blocks
1595
1596 for i, block_entry in enumerate(blocks_data):
1597 block_start_addr, block_len, tp_rank = block_entry
@@ -2041,6 +2138,17 @@ def test_compatibility_hash_validation(
2138 )
2139 decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2140 decode_worker = decode_connector.connector_worker
2141 + kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
2142 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
2143 + )
2144 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2145 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2146 + kv_caches = {
2147 + "layer0": shared_tensor,
2148 + "layer1": unique_tensor,
2149 + "layer2": shared_tensor,
2150 + }
2151 + decode_connector.register_kv_caches(kv_caches)
2152
2153 remote_config_params: dict[str, Any] = {
2154 "model": "facebook/opt-125m",
@@ -2063,7 +2171,9 @@ def test_compatibility_hash_validation(
2171 )
2172 )
2173 remote_hash = compute_nixl_compatibility_hash(
2066 - remote_vllm_config, decode_worker.backend_name
2174 + remote_vllm_config,
2175 + decode_worker.backend_name,
2176 + decode_worker.kv_topo.cross_layers_blocks,
2177 )
2178
2179 prefill_block_size = config_overrides.get("block_size", 16)
@@ -2142,6 +2252,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2252 decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2253 decode_worker = decode_connector.connector_worker
2254
2255 + backend = get_current_attn_backend(local_vllm_config)
2256 + test_shape = backend.get_kv_cache_shape(
2257 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
2258 + )
2259 + decode_worker.kv_topo = TpKVTopology(
2260 + tp_rank=decode_worker.tp_rank,
2261 + engine_id=decode_worker.engine_id,
2262 + remote_tp_size=decode_worker._tp_size, # shared state
2263 + remote_block_size=decode_worker._block_size, # shared state
2264 + is_mla=decode_worker.use_mla,
2265 + total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
2266 + attn_backend=backend,
2267 + tensor_shape=test_shape,
2268 + )
2269 +
2270 + decode_worker.compat_hash = compute_nixl_compatibility_hash(
2271 + decode_worker.vllm_config,
2272 + decode_worker.backend_name,
2273 + decode_worker.kv_topo.cross_layers_blocks,
2274 + )
2275 +
2276 if error_scenario == "handshake_decode_error":
2277 msg_bytes = b"this is not valid msgpack data"
2278 elif error_scenario == "handshake_validation_error":
vllm/distributed/kv_transfer/kv_connector/utils.py
+39 -2
@@ -316,12 +316,13 @@ class TpKVTopology:
316 attn_backend: type[AttentionBackend]
317 engine_id: EngineId
318 remote_block_size: dict[EngineId, int]
319 + tensor_shape: torch.Size | None = None
320
321 def __post_init__(self):
322 # Figure out whether the first dimension of the cache is K/V
323 # or num_blocks. This is used to register the memory regions correctly.
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=16, num_kv_heads=4, 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,32 @@ class TpKVTopology:
330 len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
331 )
332
333 + self._kv_heads_position: int | None = None
334 + self._cross_layers_blocks = False
335 + if self.tensor_shape is not None:
336 + self._cross_layers_blocks = (
337 + len(self.tensor_shape) == len(kv_cache_shape) + 1
338 + )
339 +
340 + if self._cross_layers_blocks:
341 + # prepend layers dimension
342 + kv_cache_shape = (80,) + 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 + # permute kv_cache_shape according to stride_order
351 + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
352 +
353 + physical_block_size_position = kv_cache_shape.index(16)
354 + assert physical_block_size_position is not None
355 + self._physical_block_size_position = -(
356 + len(kv_cache_shape) - physical_block_size_position
357 + )
358 +
359 @property
360 def is_kv_layout_blocks_first(self) -> bool:
361 return self._is_kv_layout_blocks_first
@@ -336,7 +363,9 @@ class TpKVTopology:
363 @property
364 def split_k_and_v(self) -> bool:
365 # Whether to register regions for K and V separately (when present).
339 - return not (self.is_mla or self.is_kv_layout_blocks_first)
366 + return not (
367 + self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
368 + )
369
370 @property
371 def tp_size(self) -> int:
@@ -346,6 +375,14 @@ class TpKVTopology:
375 def block_size(self) -> int:
376 return self.remote_block_size[self.engine_id]
377
378 + @property
379 + def cross_layers_blocks(self) -> bool:
380 + return self._cross_layers_blocks
381 +
382 + @property
383 + def block_size_position(self) -> int:
384 + return self._physical_block_size_position
385 +
386 def tp_ratio(
387 self,
388 remote_tp_size: int,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+73 -38
@@ -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,20 @@ 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 + # For now there is no benefit to run cross layers when backend
310 + # does not support on HND
311 + return False
312 +
313 + extra_config = self.kv_transfer_config.kv_connector_extra_config
314 + return bool(str(extra_config.get("enable_cross_layers_blocks", "False")))
315 +
316 def __init__(
317 self,
318 vllm_config: VllmConfig,
@@ -309,6 +324,7 @@ class NixlConnector(KVConnectorBase_V1):
324 assert vllm_config.kv_transfer_config is not None
325 assert vllm_config.kv_transfer_config.engine_id is not None
326 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
327 + self.kv_transfer_config = vllm_config.kv_transfer_config
328
329 if role == KVConnectorRole.SCHEDULER:
330 self.connector_scheduler: NixlConnectorScheduler | None = (
@@ -395,6 +411,16 @@ class NixlConnector(KVConnectorBase_V1):
411 assert self.connector_worker is not None
412 self.connector_worker.register_kv_caches(kv_caches)
413
414 + def register_cross_layers_kv_cache(
415 + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
416 + ):
417 + assert self.connector_worker is not None
418 +
419 + cross_layer_name = "ALL_LAYERS"
420 + kv_caches = {cross_layer_name: kv_cache}
421 +
422 + self.connector_worker.register_kv_caches(kv_caches)
423 +
424 def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
425 assert self.connector_worker is not None
426 self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
@@ -976,20 +1002,17 @@ class NixlConnectorWorker:
1002
1003 # Get the attention backend from the first layer
1004 # NOTE (NickLucche) models with multiple backends are not supported yet
979 - backend = get_current_attn_backend(vllm_config)
1005 + self.attn_backend = get_current_attn_backend(vllm_config)
1006
981 - self.backend_name = backend.get_name()
1007 + self.backend_name = self.attn_backend.get_name()
1008 self.kv_cache_layout = get_kv_cache_layout()
1009 self.host_buffer_kv_cache_layout = self.kv_cache_layout
1010 logger.debug("Detected attention backend %s", self.backend_name)
1011 logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
1012
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 - )
1013 + # lazy initialized in register_kv_caches
1014 + self.compat_hash: str | None = None
1015 + self.kv_topo: TpKVTopology | None = None
1016
1017 self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
1018 self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
@@ -998,17 +1021,12 @@ class NixlConnectorWorker:
1021 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
1022 self.xfer_stats = NixlKVConnectorStats()
1023
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 - )
1024 self._physical_blocks_per_logical_kv_block = 1
1025
1026 + self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
1027 + "enforce_handshake_compat", True
1028 + )
1029 +
1030 def _nixl_handshake(
1031 self,
1032 host: str,
@@ -1022,6 +1040,7 @@ class NixlConnectorWorker:
1040 # Regardless, only handshake with the remote TP rank(s) that current
1041 # local rank will read from. Note that With homogeneous TP,
1042 # this happens to be the same single rank_i.
1043 + assert self.kv_topo is not None
1044 p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
1045 remote_rank_to_agent_name = {}
1046 path = make_zmq_path("tcp", host, port)
@@ -1059,6 +1078,7 @@ class NixlConnectorWorker:
1078 )
1079
1080 # Check compatibility hash BEFORE decoding agent metadata
1081 + assert self.compat_hash is not None
1082 if (
1083 self.enforce_compat_hash
1084 and handshake_payload.compatibility_hash != self.compat_hash
@@ -1267,6 +1287,20 @@ class NixlConnectorWorker:
1287 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1288 """Register the KV Cache data in nixl."""
1289
1290 + self.kv_topo = TpKVTopology(
1291 + tp_rank=self.tp_rank,
1292 + engine_id=self.engine_id,
1293 + remote_tp_size=self._tp_size, # shared state
1294 + remote_block_size=self._block_size, # shared state
1295 + is_mla=self.use_mla,
1296 + total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1297 + attn_backend=self.attn_backend,
1298 + tensor_shape=next(iter(kv_caches.values())).shape,
1299 + )
1300 + self.compat_hash = compute_nixl_compatibility_hash(
1301 + self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
1302 + )
1303 +
1304 if self.use_host_buffer:
1305 self.initialize_host_xfer_buffer(kv_caches=kv_caches)
1306 assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -1301,29 +1335,21 @@ class NixlConnectorWorker:
1335 # (roughly 8KB vs 5KB).
1336 # Conversely for FlashInfer, K and V are registered in the same region
1337 # 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
1338 tensor_size_bytes = None
1339
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 -
1340 # Enable different block lengths for different layers when MLA is used.
1341 self.block_len_per_layer = list[int]()
1342 self.slot_size_per_layer = list[int]() # HD bytes in kv terms
1343 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 -
1344 + cache_list = (
1345 + cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
1346 + )
1347 for cache in cache_list:
1348 base_addr = cache.data_ptr()
1349 if base_addr in seen_base_addresses:
1350 continue
1351
1325 - kernel_block_size = cache.shape[block_size_position]
1326 -
1352 + kernel_block_size = cache.shape[self.kv_topo.block_size_position]
1353 if self.block_size != kernel_block_size:
1354 logger.info_once(
1355 "User-specified logical block size (%s) does not match"
@@ -1385,6 +1411,7 @@ class NixlConnectorWorker:
1411
1412 self.device_kv_caches = kv_caches
1413 self.dst_num_blocks[self.engine_id] = self.num_blocks
1414 +
1415 if self.kv_topo.is_kv_layout_blocks_first:
1416 for i in range(len(self.slot_size_per_layer)):
1417 assert self.slot_size_per_layer[i] % 2 == 0
@@ -1440,6 +1467,7 @@ class NixlConnectorWorker:
1467 block_size=self.block_size,
1468 )
1469 # Wrap metadata in payload with hash for defensive decoding
1470 + assert self.compat_hash is not None
1471 encoder = msgspec.msgpack.Encoder()
1472 self.xfer_handshake_metadata = NixlHandshakePayload(
1473 compatibility_hash=self.compat_hash,
@@ -1461,6 +1489,8 @@ class NixlConnectorWorker:
1489 register another local_xfer_handler using remote block len to ensure
1490 data copy correctness.
1491 """
1492 + assert self.kv_topo is not None
1493 +
1494 block_size_ratio = self.block_size // block_size
1495 blocks_data = []
1496 for i, base_addr in enumerate(self.seen_base_addresses):
@@ -1573,6 +1603,7 @@ class NixlConnectorWorker:
1603 # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
1604 # local origin:| 0| 1| 8| 12|
1605 # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
1606 + assert self.kv_topo is not None
1607 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
1608
1609 if engine_id not in self.dst_num_blocks:
@@ -1700,7 +1731,10 @@ class NixlConnectorWorker:
1731 """
1732 remote_engine_id = nixl_agent_meta.engine_id
1733
1703 - assert self._tp_size[remote_engine_id] == remote_tp_size
1734 + assert (
1735 + self._tp_size[remote_engine_id] == remote_tp_size
1736 + and self.kv_topo is not None
1737 + )
1738
1739 tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
1740 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
@@ -1837,6 +1871,7 @@ class NixlConnectorWorker:
1871 if len(self.device_kv_caches) == 0:
1872 return
1873 assert block_size_ratio >= 1, "Only nP < nD supported currently."
1874 + assert self.kv_topo is not None
1875 if self.enable_permute_local_kv and block_size_ratio > 1:
1876 logger.debug(
1877 "Post-processing device kv cache on receive by converting "
@@ -1856,7 +1891,7 @@ class NixlConnectorWorker:
1891 block_size_ratio,
1892 )
1893
1859 - split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
1894 + split_k_and_v = self.kv_topo.split_k_and_v
1895
1896 for block_ids in block_ids_list:
1897 indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
@@ -1881,6 +1916,7 @@ class NixlConnectorWorker:
1916 The scheduler process (via the MultiprocExecutor) will use this output
1917 to track which workers are done.
1918 """
1919 + assert self.kv_topo is not None
1920 done_sending = self._get_new_notifs()
1921 done_recving = self._pop_done_transfers(self._recving_transfers)
1922
@@ -1950,6 +1986,7 @@ class NixlConnectorWorker:
1986 are reading from the same producer (heterogeneous TP scenario), wait
1987 for all consumers to be done pulling.
1988 """
1989 + assert self.kv_topo is not None
1990 notified_req_ids: set[str] = set()
1991 for notifs in self.nixl_wrapper.get_new_notifs().values():
1992 for notif in notifs:
@@ -2109,7 +2146,7 @@ class NixlConnectorWorker:
2146 self._reqs_to_send[req_id] = expiration_time
2147
2148 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2112 - assert meta.remote is not None
2149 + assert meta.remote is not None and self.kv_topo is not None
2150 remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
2151 meta.remote.engine_id
2152 )
@@ -2178,10 +2215,7 @@ class NixlConnectorWorker:
2215 local_xfer_side_handle: int,
2216 remote_xfer_side_handle: int,
2217 ):
2181 - """
2182 - Post a READ point-to-point xfer request from a single local worker to
2183 - a single remote worker.
2184 - """
2218 + assert self.kv_topo is not None
2219 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
2220 if block_size_ratio > 1:
2221 local_block_ids = self.get_mapped_blocks(
@@ -2414,6 +2448,7 @@ class NixlConnectorWorker:
2448 For FlashInfer, this is half the length of the whole block, as K and V
2449 share the same region.
2450 """
2451 + assert self.kv_topo is not None
2452 if self.kv_topo.is_kv_layout_blocks_first:
2453 # For indexing only half (either just the K or V part).
2454 block_len = self.block_len_per_layer[layer_idx] // 2