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, 11:15 UTC
Revert "Enable Cross layers KV cache layout at NIXL Connector (#30207)" (#33241)
Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Or Ozeri committed
Jan 28, 2026 at 14:36 UTC
2e8de86777314d276ef110179a1e041c67e31e4f
5 files changed
+88
-307
docs/features/nixl_connector_usage.md
-9
@@ -184,15 +184,6 @@ 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
-
187
## Example Scripts/Code
188
189
Refer to these example scripts in the vLLM repository:
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+2
-9
@@ -34,18 +34,11 @@ 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
-
37
# Build the kv-transfer-config once
38
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
46
- KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
39
+ KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
40
else
48
- KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
41
+ KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
42
fi
43
44
# Models to run
tests/v1/kv_connector/unit/test_nixl_connector.py
+47
-178
@@ -18,12 +18,8 @@ import ray
18
import torch
19
20
from vllm import LLM
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
-)
21
+from vllm.config import KVTransferConfig
22
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
23
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
24
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
25
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
@@ -52,11 +48,8 @@ from vllm.sampling_params import SamplingParams
48
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
49
from vllm.v1.engine import EngineCoreRequest
50
from vllm.v1.engine.output_processor import OutputProcessor
55
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
51
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
52
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
53
54
from .utils import create_request, create_scheduler, create_vllm_config
55
@@ -373,7 +366,6 @@ def test_kv_transfer_handshake(dist_init):
366
367
# Decode connector will be able to create handshake with the prefill connector.
368
decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
376
- decode_connector.register_kv_caches(kv_caches)
369
370
# Here we are testing the retrieval of NIXLAgentMetadata.
371
# Knowing the implementation detail, we override the add_remote_agent
@@ -410,23 +402,6 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
402
self.kv_cache_layout = kv_cache_layout
403
# Mock register_kv_caches attribute needed for tests that do not call it.
404
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
- )
405
406
def _nixl_handshake(
407
self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
@@ -1395,7 +1370,6 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1370
),
1371
),
1372
"TRITON_ATTN",
1398
- "FLASHINFER",
1373
],
1374
)
1375
def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
@@ -1412,11 +1386,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1386
1387
vllm_config = create_vllm_config(attention_backend=attn_backend)
1388
1415
- # Enable cross layers blocks
1416
- vllm_config.kv_transfer_config.kv_connector_extra_config[
1417
- "enable_cross_layers_blocks"
1418
- ] = True
1419
-
1389
# Import the appropriate backend based on the parameter
1390
if attn_backend == "FLASH_ATTN":
1391
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
@@ -1426,11 +1395,49 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1395
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
1396
1397
backend_cls = RocmAttentionBackend
1429
- else: # TRITON
1398
+ else: # TRITON_ATTN
1399
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
1400
1401
backend_cls = TritonAttentionBackend
1402
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
+
1441
nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1442
with (
1443
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
@@ -1459,107 +1466,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1466
# Reassure the shutdown() check that the thread is terminated
1467
mock_thread.return_value.is_alive.return_value = False
1468
1462
- expected_tensor_size: int
1463
- expected_base_addrs: list[int]
1464
- expected_num_entries: int
1465
- kv_caches: dict[str, torch.Tensor]
1466
- if connector.prefer_cross_layer_blocks:
1467
- num_layers = 32
1468
- block_size = 16
1469
- num_blocks = 8
1470
- kv_cache_spec = AttentionSpec(
1471
- block_size=block_size,
1472
- num_kv_heads=4,
1473
- head_size=64,
1474
- dtype=torch.bfloat16,
1475
- )
1476
- kv_cache_config = KVCacheConfig(
1477
- num_blocks=num_blocks,
1478
- kv_cache_tensors=[
1479
- KVCacheTensor(
1480
- size=kv_cache_spec.page_size_bytes * num_blocks,
1481
- shared_by=["dummy-layer"],
1482
- )
1483
- for i in range(num_layers)
1484
- ],
1485
- # allocate_uniform_kv_caches does not use this
1486
- kv_cache_groups=[],
1487
- )
1488
-
1489
- with set_current_vllm_config(vllm_config):
1490
- _, cross_layers_kv_cache, _ = (
1491
- KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
1492
- kv_cache_config=kv_cache_config,
1493
- attn_groups=[
1494
- [
1495
- AttentionGroup(
1496
- backend=backend_cls,
1497
- layer_names=[],
1498
- kv_cache_spec=kv_cache_spec,
1499
- kv_cache_group_id=0,
1500
- )
1501
- ]
1502
- ],
1503
- cache_dtype=torch.bfloat16,
1504
- device=torch.cuda.current_device(),
1505
- kernel_block_sizes=[block_size],
1506
- )
1507
- )
1508
- # Store tensor info for validation
1509
- expected_tensor_size = (
1510
- cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
1511
- )
1512
- expected_base_addrs = [
1513
- cross_layers_kv_cache.data_ptr(),
1514
- ]
1515
- expected_num_entries = 1
1516
-
1517
- expected_blocks_count = 8
1518
-
1519
- kv_caches = {"all-layers": cross_layers_kv_cache}
1520
-
1521
- else:
1522
- # Create test kv cache tensors using proper backend shape
1523
- kv_cache_shape = backend_cls.get_kv_cache_shape(
1524
- num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
1525
- )
1526
- shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1527
- unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1528
- kv_caches = {
1529
- "layer0": shared_tensor,
1530
- "layer1": unique_tensor,
1531
- "layer2": shared_tensor,
1532
- }
1533
-
1534
- # Store tensor info for validation
1535
-
1536
- test_shape = backend_cls.get_kv_cache_shape(
1537
- num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1538
- )
1539
- is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1540
-
1541
- if is_blocks_first:
1542
- expected_tensor_size = (
1543
- shared_tensor.element_size() * shared_tensor.numel()
1544
- )
1545
- expected_base_addrs = [
1546
- shared_tensor.data_ptr(),
1547
- unique_tensor.data_ptr(),
1548
- ]
1549
- expected_num_entries = 2
1550
- else:
1551
- expected_tensor_size = (
1552
- shared_tensor[0].element_size() * shared_tensor[0].numel()
1553
- )
1554
- expected_base_addrs = [
1555
- shared_tensor[0].data_ptr(),
1556
- shared_tensor[1].data_ptr(),
1557
- unique_tensor[0].data_ptr(),
1558
- unique_tensor[1].data_ptr(),
1559
- ]
1560
- expected_num_entries = 4
1561
- expected_blocks_count = 8
1562
-
1469
# Execute register_kv_caches
1470
connector.register_kv_caches(kv_caches)
1471
@@ -1583,19 +1489,16 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1489
blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
1490
1491
# Validate blocks_data structure and size
1492
+ expected_blocks_count = 8
1493
assert len(blocks_data) == expected_blocks_count, (
1494
f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
1495
)
1496
1590
- if connector.prefer_cross_layer_blocks:
1591
- num_blocks = 8
1592
- expected_block_len = expected_tensor_size // num_blocks
1497
+ num_blocks = 2
1498
+ if is_blocks_first:
1499
+ expected_block_len = expected_tensor_size // num_blocks // 2
1500
else:
1594
- num_blocks = 2
1595
- if is_blocks_first:
1596
- expected_block_len = expected_tensor_size // num_blocks // 2
1597
- else:
1598
- expected_block_len = expected_tensor_size // num_blocks
1501
+ expected_block_len = expected_tensor_size // num_blocks
1502
1503
for i, block_entry in enumerate(blocks_data):
1504
block_start_addr, block_len, tp_rank = block_entry
@@ -2146,17 +2049,6 @@ def test_compatibility_hash_validation(
2049
)
2050
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2051
decode_worker = decode_connector.connector_worker
2149
- kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
2150
- num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
2151
- )
2152
- shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2153
- unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
2154
- kv_caches = {
2155
- "layer0": shared_tensor,
2156
- "layer1": unique_tensor,
2157
- "layer2": shared_tensor,
2158
- }
2159
- decode_connector.register_kv_caches(kv_caches)
2052
2053
remote_config_params: dict[str, Any] = {
2054
"model": "facebook/opt-125m",
@@ -2179,9 +2071,7 @@ def test_compatibility_hash_validation(
2071
)
2072
)
2073
remote_hash = compute_nixl_compatibility_hash(
2182
- remote_vllm_config,
2183
- decode_worker.backend_name,
2184
- decode_worker.kv_topo.cross_layers_blocks,
2074
+ remote_vllm_config, decode_worker.backend_name
2075
)
2076
2077
prefill_block_size = config_overrides.get("block_size", 16)
@@ -2260,27 +2150,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2150
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2151
decode_worker = decode_connector.connector_worker
2152
2263
- backend = get_current_attn_backend(local_vllm_config)
2264
- test_shape = backend.get_kv_cache_shape(
2265
- num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
2266
- )
2267
- decode_worker.kv_topo = TpKVTopology(
2268
- tp_rank=decode_worker.tp_rank,
2269
- engine_id=decode_worker.engine_id,
2270
- remote_tp_size=decode_worker._tp_size, # shared state
2271
- remote_block_size=decode_worker._block_size, # shared state
2272
- is_mla=decode_worker.use_mla,
2273
- total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
2274
- attn_backend=backend,
2275
- tensor_shape=test_shape,
2276
- )
2277
-
2278
- decode_worker.compat_hash = compute_nixl_compatibility_hash(
2279
- decode_worker.vllm_config,
2280
- decode_worker.backend_name,
2281
- decode_worker.kv_topo.cross_layers_blocks,
2282
- )
2283
-
2153
if error_scenario == "handshake_decode_error":
2154
msg_bytes = b"this is not valid msgpack data"
2155
elif error_scenario == "handshake_validation_error":
vllm/distributed/kv_transfer/kv_connector/utils.py
+1
-38
@@ -316,7 +316,6 @@ 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
319
320
def __post_init__(self):
321
# Figure out whether the first dimension of the cache is K/V
@@ -330,32 +329,6 @@ class TpKVTopology:
329
len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
330
)
331
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
-
332
@property
333
def is_kv_layout_blocks_first(self) -> bool:
334
return self._is_kv_layout_blocks_first
@@ -363,9 +336,7 @@ class TpKVTopology:
336
@property
337
def split_k_and_v(self) -> bool:
338
# Whether to register regions for K and V separately (when present).
366
- return not (
367
- self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
368
- )
339
+ return not (self.is_mla or self.is_kv_layout_blocks_first)
340
341
@property
342
def tp_size(self) -> int:
@@ -375,14 +346,6 @@ class TpKVTopology:
346
def block_size(self) -> int:
347
return self.remote_block_size[self.engine_id]
348
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
-
349
def tp_ratio(
350
self,
351
remote_tp_size: int,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+38
-73
@@ -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 AttentionBackend, AttentionMetadata
57
+from vllm.v1.attention.backend import 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, cross_layers_blocks: bool
176
+ vllm_config: VllmConfig, attn_backend_name: str
177
) -> str:
178
"""
179
Compute compatibility hash for NIXL KV transfer.
@@ -216,7 +216,6 @@ 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,
219
}
220
221
compat_hash = hash_factors(factors)
@@ -299,20 +298,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):
298
299
300
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
-
301
def __init__(
302
self,
303
vllm_config: VllmConfig,
@@ -324,7 +309,6 @@ class NixlConnector(KVConnectorBase_V1):
309
assert vllm_config.kv_transfer_config is not None
310
assert vllm_config.kv_transfer_config.engine_id is not None
311
self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
327
- self.kv_transfer_config = vllm_config.kv_transfer_config
312
313
if role == KVConnectorRole.SCHEDULER:
314
self.connector_scheduler: NixlConnectorScheduler | None = (
@@ -411,16 +395,6 @@ class NixlConnector(KVConnectorBase_V1):
395
assert self.connector_worker is not None
396
self.connector_worker.register_kv_caches(kv_caches)
397
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
-
398
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
399
assert self.connector_worker is not None
400
self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
@@ -1002,17 +976,20 @@ class NixlConnectorWorker:
976
977
# Get the attention backend from the first layer
978
# NOTE (NickLucche) models with multiple backends are not supported yet
1005
- self.attn_backend = get_current_attn_backend(vllm_config)
979
+ backend = get_current_attn_backend(vllm_config)
980
1007
- self.backend_name = self.attn_backend.get_name()
981
+ self.backend_name = backend.get_name()
982
self.kv_cache_layout = get_kv_cache_layout()
983
self.host_buffer_kv_cache_layout = self.kv_cache_layout
984
logger.debug("Detected attention backend %s", self.backend_name)
985
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
986
1013
- # lazy initialized in register_kv_caches
1014
- self.compat_hash: str | None = None
1015
- self.kv_topo: TpKVTopology | None = None
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
+ )
993
994
self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
995
self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
@@ -1021,11 +998,16 @@ class NixlConnectorWorker:
998
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
999
self.xfer_stats = NixlKVConnectorStats()
1000
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
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
)
1010
+ self._physical_blocks_per_logical_kv_block = 1
1011
1012
def _nixl_handshake(
1013
self,
@@ -1040,7 +1022,6 @@ class NixlConnectorWorker:
1022
# Regardless, only handshake with the remote TP rank(s) that current
1023
# local rank will read from. Note that With homogeneous TP,
1024
# this happens to be the same single rank_i.
1043
- assert self.kv_topo is not None
1025
p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
1026
remote_rank_to_agent_name = {}
1027
path = make_zmq_path("tcp", host, port)
@@ -1078,7 +1059,6 @@ class NixlConnectorWorker:
1059
)
1060
1061
# Check compatibility hash BEFORE decoding agent metadata
1081
- assert self.compat_hash is not None
1062
if (
1063
self.enforce_compat_hash
1064
and handshake_payload.compatibility_hash != self.compat_hash
@@ -1287,20 +1267,6 @@ class NixlConnectorWorker:
1267
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1268
"""Register the KV Cache data in nixl."""
1269
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
-
1270
if self.use_host_buffer:
1271
self.initialize_host_xfer_buffer(kv_caches=kv_caches)
1272
assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -1335,21 +1301,29 @@ class NixlConnectorWorker:
1301
# (roughly 8KB vs 5KB).
1302
# Conversely for FlashInfer, K and V are registered in the same region
1303
# 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
1305
tensor_size_bytes = None
1306
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
+
1314
# Enable different block lengths for different layers when MLA is used.
1315
self.block_len_per_layer = list[int]()
1316
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
1317
for layer_name, cache_or_caches in xfer_buffers.items():
1344
- cache_list = (
1345
- cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
1346
- )
1318
+ cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1319
+
1320
for cache in cache_list:
1321
base_addr = cache.data_ptr()
1322
if base_addr in seen_base_addresses:
1323
continue
1324
1352
- kernel_block_size = cache.shape[self.kv_topo.block_size_position]
1325
+ kernel_block_size = cache.shape[block_size_position]
1326
+
1327
if self.block_size != kernel_block_size:
1328
logger.info_once(
1329
"User-specified logical block size (%s) does not match"
@@ -1411,7 +1385,6 @@ class NixlConnectorWorker:
1385
1386
self.device_kv_caches = kv_caches
1387
self.dst_num_blocks[self.engine_id] = self.num_blocks
1414
-
1388
if self.kv_topo.is_kv_layout_blocks_first:
1389
for i in range(len(self.slot_size_per_layer)):
1390
assert self.slot_size_per_layer[i] % 2 == 0
@@ -1467,7 +1440,6 @@ class NixlConnectorWorker:
1440
block_size=self.block_size,
1441
)
1442
# Wrap metadata in payload with hash for defensive decoding
1470
- assert self.compat_hash is not None
1443
encoder = msgspec.msgpack.Encoder()
1444
self.xfer_handshake_metadata = NixlHandshakePayload(
1445
compatibility_hash=self.compat_hash,
@@ -1489,8 +1461,6 @@ class NixlConnectorWorker:
1461
register another local_xfer_handler using remote block len to ensure
1462
data copy correctness.
1463
"""
1492
- assert self.kv_topo is not None
1493
-
1464
block_size_ratio = self.block_size // block_size
1465
blocks_data = []
1466
for i, base_addr in enumerate(self.seen_base_addresses):
@@ -1603,7 +1573,6 @@ class NixlConnectorWorker:
1573
# remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
1574
# local origin:| 0| 1| 8| 12|
1575
# 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
1576
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
1577
1578
if engine_id not in self.dst_num_blocks:
@@ -1731,10 +1700,7 @@ class NixlConnectorWorker:
1700
"""
1701
remote_engine_id = nixl_agent_meta.engine_id
1702
1734
- assert (
1735
- self._tp_size[remote_engine_id] == remote_tp_size
1736
- and self.kv_topo is not None
1737
- )
1703
+ assert self._tp_size[remote_engine_id] == remote_tp_size
1704
1705
tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
1706
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
@@ -1871,7 +1837,6 @@ class NixlConnectorWorker:
1837
if len(self.device_kv_caches) == 0:
1838
return
1839
assert block_size_ratio >= 1, "Only nP < nD supported currently."
1874
- assert self.kv_topo is not None
1840
if self.enable_permute_local_kv and block_size_ratio > 1:
1841
logger.debug(
1842
"Post-processing device kv cache on receive by converting "
@@ -1891,7 +1856,7 @@ class NixlConnectorWorker:
1856
block_size_ratio,
1857
)
1858
1894
- split_k_and_v = self.kv_topo.split_k_and_v
1859
+ split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
1860
1861
for block_ids in block_ids_list:
1862
indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
@@ -1916,7 +1881,6 @@ class NixlConnectorWorker:
1881
The scheduler process (via the MultiprocExecutor) will use this output
1882
to track which workers are done.
1883
"""
1919
- assert self.kv_topo is not None
1884
done_sending = self._get_new_notifs()
1885
done_recving = self._pop_done_transfers(self._recving_transfers)
1886
@@ -1986,7 +1950,6 @@ class NixlConnectorWorker:
1950
are reading from the same producer (heterogeneous TP scenario), wait
1951
for all consumers to be done pulling.
1952
"""
1989
- assert self.kv_topo is not None
1953
notified_req_ids: set[str] = set()
1954
for notifs in self.nixl_wrapper.get_new_notifs().values():
1955
for notif in notifs:
@@ -2146,7 +2109,7 @@ class NixlConnectorWorker:
2109
self._reqs_to_send[req_id] = expiration_time
2110
2111
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2149
- assert meta.remote is not None and self.kv_topo is not None
2112
+ assert meta.remote is not None
2113
remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
2114
meta.remote.engine_id
2115
)
@@ -2215,7 +2178,10 @@ class NixlConnectorWorker:
2178
local_xfer_side_handle: int,
2179
remote_xfer_side_handle: int,
2180
):
2218
- assert self.kv_topo is not None
2181
+ """
2182
+ Post a READ point-to-point xfer request from a single local worker to
2183
+ a single remote worker.
2184
+ """
2185
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
2186
if block_size_ratio > 1:
2187
local_block_ids = self.get_mapped_blocks(
@@ -2448,7 +2414,6 @@ class NixlConnectorWorker:
2414
For FlashInfer, this is half the length of the whole block, as K and V
2415
share the same region.
2416
"""
2451
- assert self.kv_topo is not None
2417
if self.kv_topo.is_kv_layout_blocks_first:
2418
# For indexing only half (either just the K or V part).
2419
block_len = self.block_len_per_layer[layer_idx] // 2