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 5, 06:00 UTC
[KVConnector][NIXL] Organize NIXL connector into its own directory (#39354)
The number of features supported by the connector has grown substantially and the `nixl_connector.py` file has accumulated a lot of code. Creates a separate directory and isolates connector/scheduler code in the hope of improving clarity and maintainability. Further refactor of components aimed at improving clarity and simplifying code will follow soon. Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 12, 2026 at 15:10 UTC
3244a2ebf25d3896c04828b6ce073ef7848ebdc9
16 files changed
+1466
-1280
docs/features/disagg_encoder.md
+1
-1
@@ -72,4 +72,4 @@ For the PD disaggregation part, the Prefill instance receives cache exactly the
72
73
`docs/features/disagg_prefill.md` shows the brief idea about the disaggregated prefill (v0)
74
75
-We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D;
75
+We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl/` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D;
docs/mkdocs/hooks/generate_metrics.py
+1
-1
@@ -19,7 +19,7 @@ METRIC_SOURCE_FILES = [
19
"output": "spec_decode.inc.md",
20
},
21
{
22
- "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py",
22
+ "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py",
23
"output": "nixl_connector.inc.md",
24
},
25
{"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"},
tests/v1/kv_connector/unit/test_multi_connector.py
+1
-1
@@ -21,7 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
21
MultiKVConnectorStats,
22
MultiKVConnectorWorkerMetadata,
23
)
24
-from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
24
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
25
NixlKVConnectorStats,
26
)
27
from vllm.v1.kv_cache_interface import KVCacheConfig
tests/v1/kv_connector/unit/test_nixl_connector.py
+46
-45
@@ -24,13 +24,13 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
24
TpKVTopology,
25
get_current_attn_backend,
26
)
27
-from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
27
+from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
28
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
29
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
30
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
31
MultiKVConnectorStats,
32
)
32
-from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
33
- KVConnectorRole,
33
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
34
NixlAgentMetadata,
35
NixlConnector,
36
NixlConnectorMetadata,
@@ -38,6 +38,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
38
NixlConnectorWorker,
39
NixlHandshakePayload,
40
NixlKVConnectorStats,
41
+)
42
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
43
compute_nixl_compatibility_hash,
44
)
45
from vllm.distributed.kv_transfer.kv_transfer_state import (
@@ -320,7 +322,7 @@ def test_prompt_less_than_block_size():
322
323
324
@patch(
323
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
325
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
326
FakeNixlWrapper,
327
)
328
def test_kv_transfer_handshake(dist_init):
@@ -534,7 +536,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
536
537
class TestNixlHandshake:
538
@patch(
537
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
539
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
540
FakeNixlWrapper,
541
)
542
def test_multi_xfer_one_engine(
@@ -621,7 +623,7 @@ class TestNixlHandshake:
623
connector.clear_connector_metadata()
624
625
@patch(
624
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
626
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
627
FakeNixlWrapper,
628
)
629
@pytest.mark.parametrize(
@@ -691,7 +693,7 @@ class TestNixlHandshake:
693
raise TimeoutError("Took too long to complete async handshake.")
694
695
@patch(
694
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
696
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
697
FakeNixlWrapper,
698
)
699
@pytest.mark.parametrize("local_tp_size", [1, 2])
@@ -703,7 +705,7 @@ class TestNixlHandshake:
705
remote configurations.
706
"""
707
monkeypatch.setattr(
706
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size",
708
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size",
709
lambda: local_tp_size,
710
)
711
@@ -760,7 +762,7 @@ class TestNixlHandshake:
762
check_handshake(6)
763
764
@patch(
763
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
765
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
766
FakeNixlWrapper,
767
)
768
def test_prefill_tp_size_greater_than_decode_tp_size_mla(
@@ -863,7 +865,7 @@ class TestNixlHandshake:
865
assert req_id not in conn_p1.connector_worker._reqs_to_process
866
867
@patch(
866
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
868
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
869
FakeNixlWrapper,
870
)
871
def test_concurrent_load_kv(
@@ -928,7 +930,7 @@ class TestNixlHandshake:
930
raise TimeoutError("Took too long to complete async handshake.")
931
932
@patch(
931
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
933
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
934
FakeNixlWrapper,
935
)
936
def test_handshake_fails_on_kv_cache_layout_mismatch(
@@ -943,7 +945,7 @@ class TestNixlHandshake:
945
# Mock TP world size to 2 to force heterogeneous TP when
946
# remote_tp_size=1
947
with patch(
946
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501
948
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501
949
return_value=2,
950
):
951
# Initialize connector and worker (with fake NIXL wrapper)
@@ -982,7 +984,7 @@ class TestNixlHandshake:
984
worker.add_remote_agent(meta, remote_tp_size=1)
985
986
@patch(
985
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
987
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
988
FakeNixlWrapper,
989
)
990
def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
@@ -997,7 +999,7 @@ class TestNixlHandshake:
999
# Mock TP world size to 2 to force heterogeneous TP when
1000
# remote_tp_size=1
1001
with patch(
1000
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501
1002
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501
1003
return_value=2,
1004
):
1005
# Initialize connector and worker (with fake NIXL wrapper)
@@ -1042,7 +1044,7 @@ class TestNixlHandshake:
1044
# we put here is important. First run ray, it will clean up the resources, then
1045
# the rest of the tests.
1046
@patch(
1045
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1047
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1048
FakeNixlWrapper,
1049
)
1050
def test_kv_connector_stats(default_vllm_config, dist_init):
@@ -1227,8 +1229,8 @@ def test_multi_kv_connector_stats_aggregation():
1229
worker_patterns = [(2, 1), (3, 0), (0, 5)] # (Nixl, Foo)
1230
1231
worker_outputs: list[ModelRunnerOutput] = []
1230
- for i, (nixl, foo) in enumerate(worker_patterns):
1231
- stats = make_multi_stats(nixl, foo)
1232
+ for i, (nixl_count, foo) in enumerate(worker_patterns):
1233
+ stats = make_multi_stats(nixl_count, foo)
1234
output = ModelRunnerOutput(
1235
req_ids=[f"req_{i}"],
1236
req_id_to_index={f"req_{i}": 0},
@@ -1256,7 +1258,7 @@ def test_multi_kv_connector_stats_aggregation():
1258
1259
1260
@patch(
1259
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1261
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1262
FakeNixlWrapper,
1263
)
1264
def test_scheduler_kv_connector_stats_aggregation():
@@ -1324,7 +1326,7 @@ def test_scheduler_kv_connector_stats_aggregation():
1326
1327
@pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
1328
@patch(
1327
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1329
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1330
FakeNixlWrapper,
1331
)
1332
def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
@@ -1513,13 +1515,14 @@ def test_register_kv_caches(
1515
1516
backend_cls = TritonAttentionBackend
1517
1516
- nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1518
+ nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker"
1519
+ nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector"
1520
with (
1518
- patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
1519
- patch(f"{nixl_module}.threading.Event"),
1520
- patch(f"{nixl_module}.threading.Thread") as mock_thread,
1521
- patch(f"{nixl_module}.get_current_attn_backend") as mock_get_attn_backend,
1522
- patch(f"{nixl_module}.get_current_attn_backends") as mock_get_attn_backends,
1521
+ patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
1522
+ patch(f"{nixl_worker}.threading.Event"),
1523
+ patch(f"{nixl_worker}.threading.Thread") as mock_thread,
1524
+ patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend,
1525
+ patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends,
1526
):
1527
# Ensure get_attn_backend returns the correct value due to
1528
# _cached_get_attn_backend returning the backend from previous
@@ -1754,28 +1757,26 @@ def test_kv_buffer_to_nixl_memory_types(
1757
vllm_config = create_vllm_config()
1758
# Override the default memory types in the config
1759
vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device
1757
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
1760
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
1761
_NIXL_SUPPORTED_DEVICE,
1762
)
1763
1764
_NIXL_SUPPORTED_DEVICE.update(FakePlatform.get_nixl_supported_devices())
1765
1766
with (
1767
+ patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"),
1768
patch(
1765
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper"
1766
- ),
1767
- patch(
1768
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Event"
1769
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Event"
1770
),
1771
patch(
1771
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Thread"
1772
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Thread"
1773
),
1774
patch(
1774
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform",
1775
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.current_platform",
1776
FakePlatform,
1777
),
1778
patch(
1778
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector._NIXL_SUPPORTED_DEVICE",
1779
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils._NIXL_SUPPORTED_DEVICE",
1780
_NIXL_SUPPORTED_DEVICE,
1781
),
1782
): # noqa: E501
@@ -1790,7 +1791,7 @@ def test_kv_buffer_to_nixl_memory_types(
1791
1792
1793
@patch(
1793
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1794
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1795
FakeNixlWrapper,
1796
)
1797
def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
@@ -1855,7 +1856,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
1856
1857
1858
@patch(
1858
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1859
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1860
FakeNixlWrapper,
1861
)
1862
def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init):
@@ -1975,7 +1976,7 @@ class FailingNixlWrapper(FakeNixlWrapper):
1976
1977
1978
@patch(
1978
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1979
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1980
FailingNixlWrapper,
1981
)
1982
@pytest.mark.parametrize(
@@ -2065,10 +2066,10 @@ def test_transfer_failure_logging(
2066
slot_mapping={},
2067
)
2068
2068
- # Capture logs from the nixl_connector logger specifically
2069
+ # Capture logs from the nixl.worker logger specifically
2070
# vLLM loggers have propagate=False, so we need to capture directly
2071
nixl_logger = logging.getLogger(
2071
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
2072
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker"
2073
)
2074
captured_logs: list[logging.LogRecord] = []
2075
@@ -2130,7 +2131,7 @@ def test_transfer_failure_logging(
2131
2132
2133
@patch(
2133
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
2134
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2135
FailingNixlWrapper,
2136
)
2137
def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
@@ -2181,7 +2182,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2182
2183
2184
@patch(
2184
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
2185
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2186
FailingNixlWrapper,
2187
)
2188
def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):
@@ -2257,7 +2258,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2258
],
2259
)
2260
@patch(
2260
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
2261
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2262
FakeNixlWrapper,
2263
)
2264
def test_compatibility_hash_validation(
@@ -2328,7 +2329,7 @@ def test_compatibility_hash_validation(
2329
elif "connector_version" in version_override:
2330
stack.enter_context(
2331
patch.object(
2331
- nixl_connector,
2332
+ nixl.metadata,
2333
"NIXL_CONNECTOR_VERSION",
2334
version_override["connector_version"],
2335
)
@@ -2365,7 +2366,7 @@ def test_compatibility_hash_validation(
2366
# Patch zmq_ctx to return our mock socket
2367
with (
2368
patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
2368
- patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx,
2369
+ patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx,
2370
):
2371
mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
2372
@@ -2399,7 +2400,7 @@ def test_compatibility_hash_validation(
2400
],
2401
)
2402
@patch(
2402
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
2403
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2404
FakeNixlWrapper,
2405
)
2406
def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario):
@@ -2464,7 +2465,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2465
mock_socket.recv.return_value = msg_bytes
2466
with (
2467
patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
2467
- patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx,
2468
+ patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx,
2469
):
2470
mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
2471
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+9
-9
@@ -31,10 +31,10 @@ from .utils import (
31
(False, [0]),
32
],
33
)
34
-@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")
34
+@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
35
def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes):
36
"""Test sw_sizes is correctly computed based on SWA enabled/disabled."""
37
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
37
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
38
NixlConnectorScheduler,
39
)
40
@@ -65,7 +65,7 @@ def test_logical_to_kernel_block_ids_with_hma():
65
When HMA is enabled, the logical block size may differ from the kernel
66
block size. Each logical block maps to multiple kernel blocks.
67
"""
68
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
68
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
69
NixlConnectorWorker,
70
)
71
@@ -169,7 +169,7 @@ def test_nixl_metadata_hma_block_ids_structure():
169
Test that NixlConnectorMetadata correctly stores block IDs for multiple
170
KV cache groups when HMA is enabled.
171
"""
172
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
172
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
173
NixlConnectorMetadata,
174
)
175
@@ -211,7 +211,7 @@ def test_nixl_metadata_hma_block_ids_structure():
211
def test_get_block_descs_ids_hybrid_ssm():
212
"""Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM
213
when ratio=1 (no kernel block size mismatch)."""
214
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
214
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
215
NixlConnectorWorker,
216
)
217
@@ -247,7 +247,7 @@ def test_get_block_descs_ids_hybrid_ssm():
247
def test_get_block_descs_ids_kernel_block_mismatch():
248
"""Test _get_block_descs_ids uses different strides for FA (kernel blocks)
249
vs SSM (logical blocks) when ratio > 1."""
250
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
250
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
251
NixlConnectorWorker,
252
)
253
@@ -284,7 +284,7 @@ def test_get_block_descs_ids_kernel_block_mismatch():
284
def test_nixl_metadata_hybrid_ssm_block_ids():
285
"""Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
286
groups with different block counts (kernel mismatch active)."""
287
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
287
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
288
NixlConnectorMetadata,
289
)
290
@@ -392,7 +392,7 @@ def test_mamba_n1_p_side_truncation():
392
],
393
ids=["fa_swa_mamba", "fa_swa_only", "fa_only"],
394
)
395
-@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")
395
+@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
396
def test_has_mamba_init(
397
mock_platform,
398
swa_enabled,
@@ -401,7 +401,7 @@ def test_has_mamba_init(
401
expected_is_hma,
402
):
403
"""Test _has_mamba / _is_hma_required derived from kv_cache_groups."""
404
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
404
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
405
NixlConnectorScheduler,
406
)
407
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+1
-1
@@ -587,7 +587,7 @@ def test_cannot_recv():
587
assert_scheduler_empty(scheduler)
588
589
590
-@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")
590
+@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
591
def test_p_side_chunked_prefill_mamba(mock_platform):
592
"""P-side integration: Mamba N-1 truncation + chunked prefill completes.
593
tests/v1/kv_connector/unit/utils.py
+1
-1
@@ -476,7 +476,7 @@ def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):
476
477
Only sets the two flags needed by the N-1 prefill logic.
478
"""
479
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
479
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
480
NixlConnectorScheduler,
481
)
482
vllm/distributed/kv_transfer/kv_connector/factory.py
+1
-1
@@ -178,7 +178,7 @@ KVConnectorFactory.register_connector(
178
179
KVConnectorFactory.register_connector(
180
"NixlConnector",
181
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector",
181
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
182
"NixlConnector",
183
)
184
vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py
new
+31
@@ -0,0 +1,31 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NIXL KV-cache transfer connector (disaggregated prefill / decode)."""
4
+
5
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import (
6
+ NixlConnector,
7
+)
8
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
9
+ NixlAgentMetadata,
10
+ NixlConnectorMetadata,
11
+ NixlHandshakePayload,
12
+)
13
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
14
+ NixlConnectorScheduler,
15
+)
16
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
17
+ NixlKVConnectorStats,
18
+)
19
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
20
+ NixlConnectorWorker,
21
+)
22
+
23
+__all__ = [
24
+ "NixlAgentMetadata",
25
+ "NixlConnector",
26
+ "NixlConnectorMetadata",
27
+ "NixlConnectorScheduler",
28
+ "NixlConnectorWorker",
29
+ "NixlHandshakePayload",
30
+ "NixlKVConnectorStats",
31
+]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
new
+283
@@ -0,0 +1,283 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NixlConnector – thin facade that delegates to scheduler / worker."""
4
+
5
+from typing import TYPE_CHECKING, Any
6
+
7
+import torch
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.distributed.kv_transfer.kv_connector.utils import (
11
+ EngineId,
12
+ get_current_attn_backend,
13
+)
14
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15
+ CopyBlocksOp,
16
+ KVConnectorBase_V1,
17
+ KVConnectorHandshakeMetadata,
18
+ KVConnectorMetadata,
19
+ KVConnectorRole,
20
+ SupportsHMA,
21
+)
22
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
23
+ KVConnectorPromMetrics,
24
+ KVConnectorStats,
25
+ PromMetric,
26
+ PromMetricT,
27
+)
28
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
29
+ NixlConnectorMetadata,
30
+)
31
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
32
+ NixlConnectorScheduler,
33
+)
34
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
35
+ NixlKVConnectorStats,
36
+ NixlPromMetrics,
37
+)
38
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
39
+ NixlConnectorWorker,
40
+)
41
+from vllm.forward_context import ForwardContext
42
+from vllm.logger import init_logger
43
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
44
+from vllm.v1.attention.backends.utils import get_kv_cache_layout
45
+from vllm.v1.core.sched.output import SchedulerOutput
46
+from vllm.v1.kv_cache_interface import MambaSpec
47
+
48
+if TYPE_CHECKING:
49
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
50
+ from vllm.v1.kv_cache_interface import KVCacheConfig
51
+ from vllm.v1.request import Request
52
+
53
+logger = init_logger(__name__)
54
+
55
+
56
+class NixlConnector(KVConnectorBase_V1, SupportsHMA):
57
+ @property
58
+ def prefer_cross_layer_blocks(self) -> bool:
59
+ if any(
60
+ [
61
+ isinstance(group.kv_cache_spec, MambaSpec)
62
+ for group in self.kv_cache_config.kv_cache_groups
63
+ ]
64
+ ):
65
+ # Hybrid SSM models do not yet support cross-layer layout
66
+ return False
67
+
68
+ backend = get_current_attn_backend(self._vllm_config)
69
+ if backend.get_name() not in (
70
+ "FLASH_ATTN",
71
+ "FLASHINFER",
72
+ "TRITON_ATTN",
73
+ ):
74
+ return False
75
+
76
+ # For now there is no benefit to run cross layers when backend
77
+ # does not support on HND
78
+ if get_kv_cache_layout() != "HND":
79
+ return False
80
+
81
+ extra_config = self.kv_transfer_config.kv_connector_extra_config
82
+ return (
83
+ str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
84
+ == "true"
85
+ )
86
+
87
+ def __init__(
88
+ self,
89
+ vllm_config: VllmConfig,
90
+ role: KVConnectorRole,
91
+ kv_cache_config: "KVCacheConfig",
92
+ ):
93
+ super().__init__(vllm_config, role, kv_cache_config)
94
+ assert vllm_config.kv_transfer_config is not None
95
+ assert vllm_config.kv_transfer_config.engine_id is not None
96
+ self.kv_cache_config = kv_cache_config
97
+ self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
98
+ self.kv_transfer_config = vllm_config.kv_transfer_config
99
+ if role == KVConnectorRole.SCHEDULER:
100
+ self.connector_scheduler: NixlConnectorScheduler | None = (
101
+ NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
102
+ )
103
+ self.connector_worker: NixlConnectorWorker | None = None
104
+ elif role == KVConnectorRole.WORKER:
105
+ self.connector_scheduler = None
106
+ self.connector_worker = NixlConnectorWorker(
107
+ vllm_config, self.engine_id, kv_cache_config
108
+ )
109
+
110
+ ############################################################
111
+ # Class Methods
112
+ ############################################################
113
+ @classmethod
114
+ def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
115
+ if vllm_config.model_config is None:
116
+ logger.warning_once(
117
+ "Unable to detect current VLLM config. "
118
+ "Fallback to default kv cache layout."
119
+ )
120
+ return None
121
+ use_mla = vllm_config.model_config.use_mla
122
+ if use_mla:
123
+ # return None when we have mla
124
+ # as the layout should not matter in that case,
125
+ # which fallback to the default behavior.
126
+ return None
127
+ logger.info_once(
128
+ "NixlConnector setting KV cache layout to HND for better xfer performance."
129
+ )
130
+ return "HND"
131
+
132
+ ############################################################
133
+ # Scheduler Side Methods
134
+ ############################################################
135
+
136
+ def get_num_new_matched_tokens(
137
+ self, request: "Request", num_computed_tokens: int
138
+ ) -> tuple[int | None, bool]:
139
+ assert self.connector_scheduler is not None
140
+ return self.connector_scheduler.get_num_new_matched_tokens(
141
+ request, num_computed_tokens
142
+ )
143
+
144
+ def update_state_after_alloc(
145
+ self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
146
+ ):
147
+ assert self.connector_scheduler is not None
148
+ return self.connector_scheduler.update_state_after_alloc(
149
+ request, blocks, num_external_tokens
150
+ )
151
+
152
+ def build_connector_meta(
153
+ self,
154
+ scheduler_output: SchedulerOutput,
155
+ ) -> KVConnectorMetadata:
156
+ assert self.connector_scheduler is not None
157
+ return self.connector_scheduler.build_connector_meta(scheduler_output)
158
+
159
+ def request_finished(
160
+ self,
161
+ request: "Request",
162
+ block_ids: list[int],
163
+ ) -> tuple[bool, dict[str, Any] | None]:
164
+ assert self.connector_scheduler is not None
165
+ return self.connector_scheduler.request_finished(request, (block_ids,))
166
+
167
+ def request_finished_all_groups(
168
+ self,
169
+ request: "Request",
170
+ block_ids: tuple[list[int], ...],
171
+ ) -> tuple[bool, dict[str, Any] | None]:
172
+ assert self.connector_scheduler is not None
173
+ return self.connector_scheduler.request_finished(request, block_ids)
174
+
175
+ def set_xfer_handshake_metadata(
176
+ self, metadata: dict[int, KVConnectorHandshakeMetadata]
177
+ ) -> None:
178
+ """
179
+ Set the KV connector handshake metadata for this connector.
180
+
181
+ Args:
182
+ metadata (dict): the handshake metadata to set.
183
+ """
184
+ assert self.connector_scheduler is not None
185
+ self.connector_scheduler.set_xfer_handshake_metadata(metadata)
186
+
187
+ ############################################################
188
+ # Worker Side Methods
189
+ ############################################################
190
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
191
+ assert self.connector_worker is not None
192
+ self.connector_worker.register_kv_caches(kv_caches)
193
+
194
+ def register_cross_layers_kv_cache(
195
+ self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
196
+ ):
197
+ assert self.connector_worker is not None
198
+ self.connector_worker.register_cross_layers_kv_caches(kv_cache)
199
+
200
+ def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
201
+ assert self.connector_worker is not None
202
+ self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
203
+
204
+ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
205
+ """Get the finished recving and sending requests."""
206
+ assert self.connector_worker is not None
207
+ return self.connector_worker.get_finished()
208
+
209
+ def get_block_ids_with_load_errors(self) -> set[int]:
210
+ """Get block IDs that failed to load via NIXL."""
211
+ assert self.connector_worker is not None
212
+ return self.connector_worker.get_block_ids_with_load_errors()
213
+
214
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
215
+ if self.connector_worker is None:
216
+ return None
217
+ return self.connector_worker.get_kv_connector_stats()
218
+
219
+ @classmethod
220
+ def build_kv_connector_stats(
221
+ cls, data: dict[str, Any] | None = None
222
+ ) -> KVConnectorStats | None:
223
+ return (
224
+ NixlKVConnectorStats(data=data)
225
+ if data is not None
226
+ else NixlKVConnectorStats()
227
+ )
228
+
229
+ @classmethod
230
+ def build_prom_metrics(
231
+ cls,
232
+ vllm_config: VllmConfig,
233
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
234
+ labelnames: list[str],
235
+ per_engine_labelvalues: dict[int, list[object]],
236
+ ) -> KVConnectorPromMetrics:
237
+ return NixlPromMetrics(
238
+ vllm_config, metric_types, labelnames, per_engine_labelvalues
239
+ )
240
+
241
+ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
242
+ assert self.connector_worker is not None
243
+ assert isinstance(self._connector_metadata, NixlConnectorMetadata)
244
+ self.connector_worker.start_load_kv(self._connector_metadata)
245
+
246
+ def wait_for_layer_load(self, layer_name: str) -> None:
247
+ """NixlConnector does not do layerwise saving."""
248
+ pass
249
+
250
+ def save_kv_layer(
251
+ self,
252
+ layer_name: str,
253
+ kv_layer: torch.Tensor,
254
+ attn_metadata: AttentionMetadata,
255
+ **kwargs,
256
+ ) -> None:
257
+ """NixlConnector does not save explicitly."""
258
+ pass
259
+
260
+ def wait_for_save(self):
261
+ assert self.connector_worker is not None
262
+ assert isinstance(self._connector_metadata, NixlConnectorMetadata)
263
+ if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
264
+ self.connector_worker.save_kv_to_host(self._connector_metadata)
265
+
266
+ def shutdown(self):
267
+ if self.connector_worker is not None:
268
+ self.connector_worker.shutdown()
269
+ if self.connector_scheduler is not None:
270
+ self.connector_scheduler.shutdown()
271
+
272
+ def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
273
+ """
274
+ Get the KVConnector handshake metadata for this connector.
275
+ This metadata is used for out-of-band connector handshake
276
+ between P/D workers.
277
+
278
+ Returns:
279
+ KVConnectorHandshakeMetadata: the handshake metadata.
280
+ None if no handshake metadata is available.
281
+ """
282
+ assert self.connector_worker is not None
283
+ return self.connector_worker.xfer_handshake_metadata
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Metadata dataclasses and helpers for the NIXL connector."""
4
+
5
+from dataclasses import dataclass
6
+from typing import Any
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
10
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
11
+ KVConnectorHandshakeMetadata,
12
+ KVConnectorMetadata,
13
+)
14
+from vllm.logger import init_logger
15
+
16
+logger = init_logger(__name__)
17
+
18
+TransferHandle = int
19
+ReqId = str
20
+
21
+GET_META_MSG = b"get_meta_msg"
22
+#
23
+# NIXL Connector Version
24
+#
25
+# Increment this version whenever there is an incompatible change to:
26
+# - NixlAgentMetadata schema
27
+# - kv_transfer_params schema or semantics
28
+# - NIXL transfer protocol or wire format
29
+# - KV cache memory layout or block organization
30
+# - Any other change that breaks P/D interoperability
31
+#
32
+# Version History:
33
+# 1: Initial version with compatibility checking
34
+# 2: Add remote_request_id to kv_transfer_params
35
+#
36
+NIXL_CONNECTOR_VERSION: int = 2
37
+
38
+
39
+@dataclass
40
+class NixlAgentMetadata:
41
+ engine_id: str
42
+ agent_metadata: bytes
43
+ kv_caches_base_addr: list[int]
44
+ device_id: int
45
+ num_blocks: int
46
+ block_lens: list[int]
47
+ kv_cache_layout: str
48
+ block_size: int
49
+ ssm_sizes: tuple[int, int]
50
+ attn_backend_name: str
51
+
52
+
53
+@dataclass
54
+class NixlHandshakePayload(KVConnectorHandshakeMetadata):
55
+ """
56
+ Wrapper for NIXL handshake sent over the wire.
57
+
58
+ Enables two-phase decoding for graceful compatibility checking:
59
+ 1. Decode NixlHandshakePayload to get compatibility_hash
60
+ 2. Compute local hash and compare
61
+ 3. Only if hashes match, decode agent_metadata_bytes
62
+
63
+ This prevents decoder errors when NixlAgentMetadata schema is
64
+ incompatible, allowing graceful failure with clear error message.
65
+ """
66
+
67
+ compatibility_hash: str
68
+ agent_metadata_bytes: bytes # NixlAgentMetadata encoded
69
+
70
+
71
+def compute_nixl_compatibility_hash(
72
+ vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
73
+) -> str:
74
+ """
75
+ Compute compatibility hash for NIXL KV transfer.
76
+
77
+ Hash only the factors that affect whether two NIXL instances can
78
+ successfully transfer KV cache data.
79
+
80
+ Factors included:
81
+ - vLLM version and NIXL connector version
82
+ - Model architecture (name, dtype, KV heads, layers)
83
+ - KV cache format (dtype, sliding window)
84
+ - Attention backend
85
+
86
+ Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout
87
+ are validated at runtime in _validate_remote_agent_handshake and are not
88
+ included in this hash to support heterogeneous deployments.
89
+
90
+ Note - the set of factors are likely to evolve significantly over
91
+ time to be more or less permissive.
92
+
93
+ Returns:
94
+ SHA-256 hex digest
95
+ """
96
+ from vllm import __version__ as vllm_version
97
+ from vllm.config.utils import hash_factors
98
+
99
+ model_config = vllm_config.model_config
100
+ cache_config = vllm_config.cache_config
101
+ is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
102
+
103
+ factors = {
104
+ # Version compatibility
105
+ "vllm_version": vllm_version,
106
+ "nixl_connector_version": NIXL_CONNECTOR_VERSION,
107
+ # Model architecture - affects KV cache shape
108
+ "model": model_config.model,
109
+ "dtype": str(model_config.dtype),
110
+ "num_kv_heads": model_config.get_total_num_kv_heads(),
111
+ "head_size": model_config.get_head_size(),
112
+ "num_hidden_layers": model_config.get_total_num_hidden_layers(),
113
+ # Attention backend and KV cache dtype affect memory layout
114
+ "attn_backend_name": attn_backend_name,
115
+ "cache_dtype": str(cache_config.cache_dtype),
116
+ "cross_layers_blocks": cross_layers_blocks,
117
+ "is_hma_enabled": is_hma_enabled,
118
+ }
119
+
120
+ compat_hash = hash_factors(factors)
121
+ logger.debug(
122
+ "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, "
123
+ "cache_dtype=%s, attn_backend=%s)",
124
+ compat_hash,
125
+ factors["model"],
126
+ factors["dtype"],
127
+ factors["num_kv_heads"],
128
+ factors["cache_dtype"],
129
+ attn_backend_name,
130
+ )
131
+ return compat_hash
132
+
133
+
134
+@dataclass
135
+class RemoteMeta:
136
+ block_ids: BlockIds
137
+ host: str
138
+ port: int
139
+ engine_id: str
140
+ request_id: str
141
+
142
+
143
+@dataclass
144
+class ReqMeta:
145
+ local_block_ids: BlockIds
146
+ # To be used when logical block size does not match the kernel block size
147
+ local_physical_block_ids: BlockIds
148
+ tp_size: int
149
+ remote: RemoteMeta | None = None
150
+
151
+
152
+class NixlConnectorMetadata(KVConnectorMetadata):
153
+ def __init__(self):
154
+ self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
155
+ self.reqs_to_save: dict[ReqId, ReqMeta] = {}
156
+ self.reqs_to_send: dict[ReqId, float] = {}
157
+ self.reqs_in_batch: set[ReqId] = set()
158
+ self.reqs_not_processed: set[ReqId] = set()
159
+
160
+ def _add_new_req(
161
+ self,
162
+ local_block_ids: BlockIds,
163
+ kv_transfer_params: dict[str, Any],
164
+ ) -> ReqMeta:
165
+ return ReqMeta(
166
+ local_block_ids=local_block_ids,
167
+ local_physical_block_ids=local_block_ids,
168
+ # P workers don't need to receive tp_size from proxy here.
169
+ tp_size=kv_transfer_params.get("tp_size", 1),
170
+ )
171
+
172
+ def add_new_req_to_save(
173
+ self,
174
+ request_id: ReqId,
175
+ local_block_ids: BlockIds,
176
+ kv_transfer_params: dict[str, Any],
177
+ ):
178
+ self.reqs_to_save[request_id] = self._add_new_req(
179
+ local_block_ids, kv_transfer_params
180
+ )
181
+
182
+ def add_new_req_to_recv(
183
+ self,
184
+ request_id: ReqId,
185
+ local_block_ids: BlockIds,
186
+ kv_transfer_params: dict[str, Any],
187
+ ):
188
+ req = self._add_new_req(local_block_ids, kv_transfer_params)
189
+ req.remote = RemoteMeta(
190
+ block_ids=kv_transfer_params["remote_block_ids"],
191
+ engine_id=kv_transfer_params["remote_engine_id"],
192
+ request_id=kv_transfer_params["remote_request_id"],
193
+ host=kv_transfer_params["remote_host"],
194
+ port=kv_transfer_params["remote_port"],
195
+ )
196
+ self.reqs_to_recv[request_id] = req
vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py
new
+504
@@ -0,0 +1,504 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scheduler-side logic for the NIXL connector."""
4
+
5
+import threading
6
+import time
7
+from typing import TYPE_CHECKING, Any
8
+
9
+import msgspec
10
+import zmq
11
+
12
+from vllm import envs
13
+from vllm.distributed.kv_transfer.kv_connector.utils import (
14
+ BlockIds,
15
+ EngineId,
16
+ yield_req_data,
17
+)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19
+ KVConnectorHandshakeMetadata,
20
+ KVConnectorMetadata,
21
+)
22
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
23
+ GET_META_MSG,
24
+ NixlConnectorMetadata,
25
+ NixlHandshakePayload,
26
+ ReqId,
27
+)
28
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx
29
+from vllm.logger import init_logger
30
+from vllm.platforms import current_platform
31
+from vllm.utils.math_utils import cdiv
32
+from vllm.utils.network_utils import make_zmq_path
33
+from vllm.v1.core.sched.output import SchedulerOutput
34
+from vllm.v1.kv_cache_interface import (
35
+ FullAttentionSpec,
36
+ MambaSpec,
37
+ SlidingWindowSpec,
38
+)
39
+
40
+if TYPE_CHECKING:
41
+ from vllm.config import VllmConfig
42
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
43
+ from vllm.v1.kv_cache_interface import KVCacheConfig
44
+ from vllm.v1.request import Request
45
+
46
+logger = init_logger(__name__)
47
+
48
+
49
+class NixlConnectorScheduler:
50
+ """Implementation of Scheduler side methods"""
51
+
52
+ def __init__(
53
+ self,
54
+ vllm_config: "VllmConfig",
55
+ engine_id: str,
56
+ kv_cache_config: "KVCacheConfig",
57
+ ):
58
+ self.vllm_config = vllm_config
59
+ self.block_size = vllm_config.cache_config.block_size
60
+ self.engine_id: EngineId = engine_id
61
+ self.kv_cache_config = kv_cache_config
62
+ self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
63
+ self.side_channel_port = (
64
+ envs.VLLM_NIXL_SIDE_CHANNEL_PORT
65
+ + vllm_config.parallel_config.data_parallel_index
66
+ )
67
+ assert vllm_config.kv_transfer_config is not None
68
+ if current_platform.device_type == "cpu":
69
+ self.use_host_buffer = False
70
+ else:
71
+ self.use_host_buffer = (
72
+ vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
73
+ )
74
+ self._is_hma_required = (
75
+ not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
76
+ # Also handle unlikely SW-only model case instead of checking num_groups>1.
77
+ and any(
78
+ not isinstance(g.kv_cache_spec, FullAttentionSpec)
79
+ for g in kv_cache_config.kv_cache_groups
80
+ )
81
+ )
82
+ self._has_mamba = any(
83
+ isinstance(g.kv_cache_spec, MambaSpec)
84
+ for g in kv_cache_config.kv_cache_groups
85
+ )
86
+
87
+ logger.info("Initializing NIXL Scheduler %s", engine_id)
88
+ if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
89
+ logger.info("Hybrid Memory Allocator is enabled with NIXL")
90
+
91
+ # Background thread for handling new handshake requests.
92
+ self._nixl_handshake_listener_t: threading.Thread | None = None
93
+ self._stop_event = threading.Event()
94
+
95
+ # Requests that need to start recv/send.
96
+ # New requests are added by update_state_after_alloc in
97
+ # the scheduler. Used to make metadata passed to Worker.
98
+ self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
99
+ self._reqs_need_save: dict[ReqId, Request] = {}
100
+ # Reqs to send and their expiration time
101
+ self._reqs_need_send: dict[ReqId, float] = {}
102
+ self._reqs_in_batch: set[ReqId] = set()
103
+ # Reqs to remove from processed set because they're not to send after
104
+ # remote prefill or aborted.
105
+ self._reqs_not_processed: set[ReqId] = set()
106
+
107
+ # Gather Sliding Window sizes for each kv cache group (if any) in number of
108
+ # blocks per KV cache group. This is used to clip the local attention window.
109
+ sw_sizes_tokens: list[tuple[int, int]] = [
110
+ (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
111
+ if isinstance(g.kv_cache_spec, SlidingWindowSpec)
112
+ else (0, self.block_size)
113
+ for g in kv_cache_config.kv_cache_groups
114
+ ]
115
+ # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
116
+ # account for boundary overlap eg window isn't fully aligned with blocks.
117
+ self.blocks_per_sw = [
118
+ cdiv(n_tokens, block_size) + 1 if n_tokens else 0
119
+ for n_tokens, block_size in sw_sizes_tokens
120
+ ]
121
+
122
+ def shutdown(self):
123
+ self._stop_event.set()
124
+ if self._nixl_handshake_listener_t is not None:
125
+ self._nixl_handshake_listener_t.join()
126
+ self._nixl_handshake_listener_t = None
127
+
128
+ def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
129
+ """
130
+ Clip the number of blocks to the sliding window size for each kv cache group
131
+ that employs SWA.
132
+ This is necessary because the KV Cache manager initially allocates blocks for
133
+ the entire sequence length, and successively cleans up blocks that are outside
134
+ the window prior to the `request_finished_all_groups` hook.
135
+ """
136
+ if len(block_ids) == 0 or not self._is_hma_required:
137
+ # No blocks to clip eg Full prefix cache hit or not a hybrid model.
138
+ return block_ids
139
+ # NOTE (NickLucche) This logic is currently handled at the connector level
140
+ # because offloading connectors might want to receive the whole sequence even
141
+ # for SWA groups. We will abstract this logic once the interface is more stable
142
+ assert len(block_ids) == len(self.blocks_per_sw), (
143
+ "Number of KV cache groups must match"
144
+ )
145
+ # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
146
+ return tuple(
147
+ [
148
+ blocks[-self.blocks_per_sw[i] :]
149
+ if self.blocks_per_sw[i] > 0
150
+ else blocks
151
+ for i, blocks in enumerate(block_ids)
152
+ ]
153
+ )
154
+
155
+ def set_xfer_handshake_metadata(
156
+ self, metadata: dict[int, KVConnectorHandshakeMetadata]
157
+ ) -> None:
158
+ """
159
+ Set the KV connector handshake metadata for this connector.
160
+
161
+ Args:
162
+ metadata (dict): the handshake metadata to set.
163
+ """
164
+ encoded_data: dict[int, bytes] = {}
165
+ encoder = msgspec.msgpack.Encoder()
166
+ for tp_rank, rank_metadata in metadata.items():
167
+ if not isinstance(rank_metadata, NixlHandshakePayload):
168
+ raise ValueError(
169
+ "NixlConnectorScheduler expects NixlHandshakePayload for "
170
+ "handshake metadata."
171
+ )
172
+ encoded_data[tp_rank] = encoder.encode(rank_metadata)
173
+ logger.debug(
174
+ "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",
175
+ tp_rank,
176
+ str(len(encoded_data[tp_rank])),
177
+ )
178
+
179
+ # Only start the listener when we have metadata to serve.
180
+ if self._nixl_handshake_listener_t is None:
181
+ ready_event = threading.Event()
182
+ self._nixl_handshake_listener_t = threading.Thread(
183
+ target=self._nixl_handshake_listener,
184
+ args=(
185
+ encoded_data,
186
+ ready_event,
187
+ self._stop_event,
188
+ self.side_channel_port,
189
+ ),
190
+ daemon=True,
191
+ name="nixl_handshake_listener",
192
+ )
193
+ self._nixl_handshake_listener_t.start()
194
+ ready_event.wait() # Wait for listener ZMQ socket to be ready.
195
+
196
+ @staticmethod
197
+ def _nixl_handshake_listener(
198
+ encoded_data: dict[int, Any],
199
+ ready_event: threading.Event,
200
+ stop_event: threading.Event,
201
+ port: int,
202
+ ):
203
+ """Background thread for getting new NIXL handshakes."""
204
+ # NOTE(rob): this is a simple implementation. We will move
205
+ # to a better approach via HTTP endpoint soon.
206
+
207
+ # Listen for new requests for metadata.
208
+ host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
209
+ path = make_zmq_path("tcp", host, port)
210
+ logger.debug("Starting listening on path: %s", path)
211
+ with zmq_ctx(zmq.ROUTER, path) as sock:
212
+ sock.setsockopt(zmq.RCVTIMEO, 1000)
213
+ ready_event.set()
214
+ while True:
215
+ try:
216
+ identity, _, msg = sock.recv_multipart()
217
+ except zmq.Again:
218
+ if stop_event.is_set():
219
+ break
220
+ continue
221
+ # Decode the message which contains (GET_META_MSG, rank)
222
+ msg, target_tp_rank = msgspec.msgpack.decode(msg)
223
+ logger.debug(
224
+ "Received message for tp rank %s",
225
+ target_tp_rank,
226
+ )
227
+ if msg != GET_META_MSG:
228
+ logger.warning("Connection listener got unexpected message %s", msg)
229
+ sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
230
+
231
+ def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
232
+ """D-side only. Returns N-1 for Mamba models since the decoder
233
+ always recomputes the last token and must start from h(N-1)."""
234
+ if self._has_mamba and num_prompt_tokens > 1:
235
+ return num_prompt_tokens - 1
236
+ return num_prompt_tokens
237
+
238
+ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
239
+ """P-side only: drop the last prompt token so the prefiller computes
240
+ h(N-1) instead of h(N). The decoder recomputes the last token to
241
+ derive h(N) correctly.
242
+
243
+ Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
244
+ request is preempted and rescheduled."""
245
+ params = request.kv_transfer_params
246
+ if (
247
+ params is not None
248
+ # Guard against repeated truncation after preemption/reschedule.
249
+ and not params.get("_p_side_truncated")
250
+ and request.num_prompt_tokens > 1
251
+ ):
252
+ if request.prompt_token_ids is not None:
253
+ request.prompt_token_ids.pop()
254
+ elif request.prompt_embeds is not None:
255
+ request.prompt_embeds = request.prompt_embeds[:-1]
256
+ else:
257
+ return
258
+
259
+ request._all_token_ids.pop()
260
+ request.num_prompt_tokens -= 1
261
+ request.max_tokens = 1
262
+ params["_p_side_truncated"] = True
263
+
264
+ def get_num_new_matched_tokens(
265
+ self, request: "Request", num_computed_tokens: int
266
+ ) -> tuple[int, bool]:
267
+ """
268
+ For remote prefill, pull all prompt blocks from remote
269
+ asynchronously relative to engine execution.
270
+
271
+ Args:
272
+ request (Request): the request object.
273
+ num_computed_tokens (int): the number of locally
274
+ computed tokens for this request
275
+ Returns:
276
+ * the number of tokens that can be loaded from the
277
+ external KV cache beyond what is already computed.
278
+ * true if the external KV cache tokens will be loaded
279
+ asynchronously (between scheduler steps).
280
+ """
281
+
282
+ params = request.kv_transfer_params
283
+ logger.debug(
284
+ "NIXLConnector get_num_new_matched_tokens: "
285
+ "num_computed_tokens=%s, kv_transfer_params=%s",
286
+ num_computed_tokens,
287
+ params,
288
+ )
289
+
290
+ if params is not None and params.get("do_remote_prefill"):
291
+ # Remote prefill: get all prompt blocks from remote.
292
+ token_ids = request.prompt_token_ids or []
293
+ actual = self._mamba_prefill_token_count(len(token_ids))
294
+ count = actual - num_computed_tokens
295
+ if count > 0:
296
+ return count, True
297
+
298
+ if params is not None and params.get("do_remote_decode") and self._has_mamba:
299
+ self._truncate_mamba_request_for_prefill(request)
300
+
301
+ # No remote prefill for this request.
302
+ return 0, False
303
+
304
+ def update_state_after_alloc(
305
+ self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
306
+ ):
307
+ params = request.kv_transfer_params
308
+ logger.debug(
309
+ "NIXLConnector update_state_after_alloc: "
310
+ "num_external_tokens=%s, kv_transfer_params=%s",
311
+ num_external_tokens,
312
+ params,
313
+ )
314
+
315
+ if not params:
316
+ return
317
+
318
+ if params.get("do_remote_decode"):
319
+ self._reqs_in_batch.add(request.request_id)
320
+ if self.use_host_buffer and params.get("do_remote_decode"):
321
+ # NOTE: when accelerator is not directly supported by Nixl,
322
+ # prefilled blocks need to be saved to host memory before transfer.
323
+ self._reqs_need_save[request.request_id] = request
324
+ elif params.get("do_remote_prefill"):
325
+ if params.get("remote_block_ids"):
326
+ if all(
327
+ p in params
328
+ for p in (
329
+ "remote_engine_id",
330
+ "remote_request_id",
331
+ "remote_host",
332
+ "remote_port",
333
+ )
334
+ ):
335
+ # If remote_blocks and num_external_tokens = 0, we have
336
+ # a full prefix cache hit on the D worker. We need to call
337
+ # send_notif in _read_blocks to free the memory on the P.
338
+
339
+ unhashed_local_block_ids: BlockIds = (
340
+ blocks.get_unhashed_block_ids_all_groups()
341
+ if num_external_tokens > 0
342
+ else ()
343
+ )
344
+ local_block_ids = self.get_sw_clipped_blocks(
345
+ unhashed_local_block_ids
346
+ )
347
+
348
+ # Get unhashed blocks to pull from remote. Mind that a full prefix
349
+ # cache hit is indicated with an empty list.
350
+ self._reqs_need_recv[request.request_id] = (
351
+ request,
352
+ local_block_ids,
353
+ )
354
+
355
+ else:
356
+ logger.warning(
357
+ "Got invalid KVTransferParams: %s. This "
358
+ "request will not utilize KVTransfer",
359
+ params,
360
+ )
361
+ else:
362
+ assert num_external_tokens == 0
363
+ # Only trigger 1 KV transfer per request.
364
+ params["do_remote_prefill"] = False
365
+
366
+ def _build_save_meta(
367
+ self,
368
+ meta: NixlConnectorMetadata,
369
+ scheduler_output: SchedulerOutput,
370
+ ) -> None:
371
+ # only called when use_host_buffer is True to build the save metadata
372
+
373
+ # NOTE: For the prefill side, there might be a chance that an early added
374
+ # request is a chunked prefill, so we need to check if new blocks are added
375
+ for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
376
+ req_to_save = self._reqs_need_save.get(req_id)
377
+ if req_to_save is None or new_block_id_groups is None:
378
+ continue
379
+ req = req_to_save
380
+
381
+ assert req.kv_transfer_params is not None
382
+ clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
383
+ meta.add_new_req_to_save(
384
+ request_id=req_id,
385
+ local_block_ids=clipped_block_id_groups,
386
+ kv_transfer_params=req.kv_transfer_params,
387
+ )
388
+ assert scheduler_output.num_scheduled_tokens is not None
389
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
390
+ is_partial = (
391
+ req.num_computed_tokens + num_scheduled_tokens
392
+ ) < req.num_prompt_tokens
393
+ if not is_partial:
394
+ # For non-partial prefills, once new req_meta is scheduled, it
395
+ # can be removed from _reqs_need_save.
396
+ # For partial prefill case, we will retain the request in
397
+ # _reqs_need_save until all blocks are scheduled with req_meta.
398
+ # Therefore, only pop if `not is_partial`.
399
+ self._reqs_need_save.pop(req_id)
400
+
401
+ def build_connector_meta(
402
+ self,
403
+ scheduler_output: SchedulerOutput,
404
+ ) -> KVConnectorMetadata:
405
+ meta = NixlConnectorMetadata()
406
+
407
+ # Loop through scheduled reqs and convert to ReqMeta.
408
+ for req_id, (req, block_ids) in self._reqs_need_recv.items():
409
+ assert req.kv_transfer_params is not None
410
+ meta.add_new_req_to_recv(
411
+ request_id=req_id,
412
+ local_block_ids=block_ids,
413
+ kv_transfer_params=req.kv_transfer_params,
414
+ )
415
+
416
+ if self.use_host_buffer:
417
+ self._build_save_meta(meta, scheduler_output)
418
+
419
+ meta.reqs_to_send = self._reqs_need_send
420
+ meta.reqs_in_batch = self._reqs_in_batch
421
+ meta.reqs_not_processed = self._reqs_not_processed
422
+
423
+ # Clear the list once workers start the transfers
424
+ self._reqs_need_recv.clear()
425
+ self._reqs_in_batch = set()
426
+ self._reqs_not_processed = set()
427
+ self._reqs_need_send = {}
428
+
429
+ return meta
430
+
431
+ def request_finished(
432
+ self,
433
+ request: "Request",
434
+ block_ids: BlockIds,
435
+ ) -> tuple[bool, dict[str, Any] | None]:
436
+ """
437
+ Once a request is finished, determine whether request blocks
438
+ should be freed now or will be sent asynchronously and freed later.
439
+ """
440
+ from vllm.v1.request import RequestStatus
441
+
442
+ params = request.kv_transfer_params
443
+ logger.debug(
444
+ "NIXLConnector request_finished(%s), request_status=%s, "
445
+ "kv_transfer_params=%s",
446
+ request.request_id,
447
+ request.status,
448
+ params,
449
+ )
450
+ if not params:
451
+ return False, None
452
+
453
+ if params.get("do_remote_prefill"):
454
+ # If do_remote_prefill is still True when the request is finished,
455
+ # update_state_after_alloc must not have been called (the request
456
+ # must have been aborted before it was scheduled).
457
+ # To avoid stranding the prefill blocks in the prefill instance,
458
+ # we must add empty block_ids to _reqs_need_recv so that our
459
+ # worker side will notify and free blocks in the prefill instance.
460
+ self._reqs_need_recv[request.request_id] = (request, [])
461
+ params["do_remote_prefill"] = False
462
+ return False, None
463
+
464
+ if not params.get("do_remote_decode"):
465
+ return False, None
466
+ if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
467
+ # Also include the case of a P/D Prefill request with immediate
468
+ # block free (eg abort). Stop tracking this request.
469
+ self._reqs_not_processed.add(request.request_id)
470
+ # Clear _reqs_need_save if a request is aborted as partial prefill.
471
+ self._reqs_need_save.pop(request.request_id, None)
472
+ return False, None
473
+
474
+ # TODO: check whether block_ids actually ever be 0. If not we could
475
+ # remove the conditional below
476
+ delay_free_blocks = any(len(group) > 0 for group in block_ids)
477
+
478
+ if delay_free_blocks:
479
+ # Prefill request on remote. It will be read from D upon completion
480
+ logger.debug(
481
+ "NIXLConnector request_finished(%s) waiting for %d seconds "
482
+ "for remote decode to fetch blocks",
483
+ request.request_id,
484
+ envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT,
485
+ )
486
+ self._reqs_need_send[request.request_id] = (
487
+ time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
488
+ )
489
+ # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
490
+ # trimming down after allocating for the whole sequence length. Empty
491
+ # blocks are always at the start of the list.
492
+ # Here we "unpad" blocks to send the actual remote blocks to be read.
493
+ block_ids = self.get_sw_clipped_blocks(block_ids)
494
+
495
+ return delay_free_blocks, dict(
496
+ do_remote_prefill=True,
497
+ do_remote_decode=False,
498
+ remote_block_ids=block_ids,
499
+ remote_engine_id=self.engine_id,
500
+ remote_request_id=request.request_id,
501
+ remote_host=self.side_channel_host,
502
+ remote_port=self.side_channel_port,
503
+ tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
504
+ )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
new
+266
@@ -0,0 +1,266 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Stats and Prometheus metrics for the NIXL connector."""
4
+
5
+import copy
6
+from dataclasses import dataclass
7
+from typing import Any
8
+
9
+import numpy as np
10
+
11
+from vllm.config import VllmConfig
12
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
13
+ KVConnectorPromMetrics,
14
+ KVConnectorStats,
15
+ PromMetric,
16
+ PromMetricT,
17
+)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
19
+ nixlXferTelemetry,
20
+)
21
+from vllm.v1.metrics.utils import create_metric_per_engine
22
+
23
+
24
+@dataclass
25
+class NixlKVConnectorStats(KVConnectorStats):
26
+ """Container for transfer performance metrics"""
27
+
28
+ def __post_init__(self):
29
+ if not self.data:
30
+ # Empty container init, no data is passed in.
31
+ self.reset()
32
+
33
+ def reset(self):
34
+ # Must be serializable
35
+ self.data: dict[str, list[float | int]] = {
36
+ "transfer_duration": [],
37
+ "post_duration": [],
38
+ "bytes_transferred": [],
39
+ "num_descriptors": [],
40
+ "num_failed_transfers": [],
41
+ "num_failed_notifications": [],
42
+ "num_kv_expired_reqs": [],
43
+ }
44
+
45
+ def record_transfer(self, res: nixlXferTelemetry):
46
+ # Keep metrics units consistent with rest of the code: time us->s
47
+ self.data["transfer_duration"].append(res.xferDuration / 1e6)
48
+ self.data["post_duration"].append(res.postDuration / 1e6)
49
+ self.data["bytes_transferred"].append(res.totalBytes)
50
+ self.data["num_descriptors"].append(res.descCount)
51
+
52
+ def record_failed_transfer(self):
53
+ """Record a failed NIXL transfer operation."""
54
+ self.data["num_failed_transfers"].append(1)
55
+
56
+ def record_failed_notification(self):
57
+ """Record a failed NIXL notification (send_notif)."""
58
+ self.data["num_failed_notifications"].append(1)
59
+
60
+ def record_kv_expired_req(self):
61
+ """Record a request that had its KV blocks expire."""
62
+ self.data["num_kv_expired_reqs"].append(1)
63
+
64
+ def clone_and_reset(self) -> "NixlKVConnectorStats":
65
+ old = copy.copy(self)
66
+ self.reset()
67
+ return old
68
+
69
+ def is_empty(self) -> bool:
70
+ # Do not discard metrics update that are entirely failures related.
71
+ return (
72
+ self.num_successful_transfers == 0
73
+ and len(self.data["num_failed_transfers"]) == 0
74
+ and len(self.data["num_failed_notifications"]) == 0
75
+ and len(self.data["num_kv_expired_reqs"]) == 0
76
+ )
77
+
78
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
79
+ if not other.is_empty():
80
+ for k, v in other.data.items():
81
+ accumulator = self.data[k]
82
+ assert isinstance(accumulator, list)
83
+ accumulator.extend(v)
84
+ return self
85
+
86
+ def reduce(self) -> dict[str, int | float]:
87
+ # Compute compact representative stats suitable for CLI logging
88
+ if self.num_successful_transfers == 0:
89
+ # CLI logging only reports successful transfers stats. If all requests in
90
+ # the interval were unsuccessful, Prom will report failures stats instead.
91
+ return {
92
+ "Num successful transfers": 0,
93
+ "Avg xfer time (ms)": 0,
94
+ "P90 xfer time (ms)": 0,
95
+ "Avg post time (ms)": 0,
96
+ "P90 post time (ms)": 0,
97
+ "Avg MB per transfer": 0,
98
+ "Throughput (MB/s)": 0,
99
+ "Avg number of descriptors": 0,
100
+ }
101
+
102
+ xfer_time = np.asarray(self.data["transfer_duration"])
103
+ post_time = np.asarray(self.data["post_duration"])
104
+ # Convert to MB for CLI logging.
105
+ mb = np.asarray(self.data["bytes_transferred"]) / 2**20
106
+ descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32)
107
+ n = len(descs)
108
+ assert n == self.num_successful_transfers
109
+
110
+ total_mb = mb.sum()
111
+ avg_mb = total_mb / n
112
+
113
+ total_time_seconds = xfer_time.sum()
114
+ throughput_mb_s = total_mb / total_time_seconds
115
+
116
+ return {
117
+ "Num successful transfers": n,
118
+ "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
119
+ "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
120
+ "Avg post time (ms)": round(post_time.mean() * 1e3, 3),
121
+ "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),
122
+ "Avg MB per transfer": round(avg_mb, 3),
123
+ "Throughput (MB/s)": round(throughput_mb_s, 3),
124
+ "Avg number of descriptors": round(descs.mean(), 1),
125
+ }
126
+
127
+ @property
128
+ def num_successful_transfers(self) -> int:
129
+ return len(self.data["transfer_duration"])
130
+
131
+
132
+class NixlPromMetrics(KVConnectorPromMetrics):
133
+ def __init__(
134
+ self,
135
+ vllm_config: VllmConfig,
136
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
137
+ labelnames: list[str],
138
+ per_engine_labelvalues: dict[int, list[object]],
139
+ ):
140
+ super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
141
+
142
+ buckets = [
143
+ 0.001,
144
+ 0.005,
145
+ 0.01,
146
+ 0.025,
147
+ 0.05,
148
+ 0.075,
149
+ 0.1,
150
+ 0.2,
151
+ 0.3,
152
+ 0.5,
153
+ 0.75,
154
+ 1.0,
155
+ 5.0,
156
+ ]
157
+ nixl_histogram_xfer_time = self._histogram_cls(
158
+ name="vllm:nixl_xfer_time_seconds",
159
+ documentation="Histogram of transfer duration for NIXL KV Cache transfers.",
160
+ buckets=buckets[1:],
161
+ labelnames=labelnames,
162
+ )
163
+ self.nixl_histogram_xfer_time = create_metric_per_engine(
164
+ nixl_histogram_xfer_time, self.per_engine_labelvalues
165
+ )
166
+ nixl_histogram_post_time = self._histogram_cls(
167
+ name="vllm:nixl_post_time_seconds",
168
+ documentation="Histogram of transfer post time for NIXL KV"
169
+ " Cache transfers.",
170
+ buckets=buckets,
171
+ labelnames=labelnames,
172
+ )
173
+ self.nixl_histogram_post_time = create_metric_per_engine(
174
+ nixl_histogram_post_time, self.per_engine_labelvalues
175
+ )
176
+ # uniform 2kb to 16gb range
177
+ buckets = [2 ** (10 + i) for i in range(1, 25, 2)]
178
+ nixl_histogram_bytes_transferred = self._histogram_cls(
179
+ name="vllm:nixl_bytes_transferred",
180
+ documentation="Histogram of bytes transferred per NIXL KV Cache transfers.",
181
+ buckets=buckets,
182
+ labelnames=labelnames,
183
+ )
184
+ self.nixl_histogram_bytes_transferred = create_metric_per_engine(
185
+ nixl_histogram_bytes_transferred, self.per_engine_labelvalues
186
+ )
187
+ buckets = [
188
+ 10,
189
+ 20,
190
+ 30,
191
+ 50,
192
+ 75,
193
+ 100,
194
+ 200,
195
+ 400,
196
+ 1000,
197
+ 2000,
198
+ 4000,
199
+ 10000,
200
+ 20000,
201
+ 50000,
202
+ ]
203
+ nixl_histogram_num_descriptors = self._histogram_cls(
204
+ name="vllm:nixl_num_descriptors",
205
+ documentation="Histogram of number of descriptors per NIXL"
206
+ " KV Cache transfers.",
207
+ buckets=buckets,
208
+ labelnames=labelnames,
209
+ )
210
+ self.nixl_histogram_num_descriptors = create_metric_per_engine(
211
+ nixl_histogram_num_descriptors, self.per_engine_labelvalues
212
+ )
213
+ counter_nixl_num_failed_transfers = self._counter_cls(
214
+ name="vllm:nixl_num_failed_transfers",
215
+ documentation="Number of failed NIXL KV Cache transfers.",
216
+ labelnames=labelnames,
217
+ )
218
+ self.counter_nixl_num_failed_transfers = create_metric_per_engine(
219
+ counter_nixl_num_failed_transfers, self.per_engine_labelvalues
220
+ )
221
+ counter_nixl_num_failed_notifications = self._counter_cls(
222
+ name="vllm:nixl_num_failed_notifications",
223
+ documentation="Number of failed NIXL KV Cache notifications.",
224
+ labelnames=labelnames,
225
+ )
226
+ self.counter_nixl_num_failed_notifications = create_metric_per_engine(
227
+ counter_nixl_num_failed_notifications, self.per_engine_labelvalues
228
+ )
229
+
230
+ counter_nixl_num_kv_expired_reqs = self._counter_cls(
231
+ name="vllm:nixl_num_kv_expired_reqs",
232
+ documentation="Number of requests that had their KV expire. "
233
+ "NOTE: This metric is tracked on the P instance.",
234
+ labelnames=labelnames,
235
+ )
236
+ self.counter_nixl_num_kv_expired_reqs = create_metric_per_engine(
237
+ counter_nixl_num_kv_expired_reqs, self.per_engine_labelvalues
238
+ )
239
+
240
+ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
241
+ for prom_obj, list_item_key in zip(
242
+ [
243
+ self.nixl_histogram_xfer_time,
244
+ self.nixl_histogram_post_time,
245
+ self.nixl_histogram_bytes_transferred,
246
+ self.nixl_histogram_num_descriptors,
247
+ ],
248
+ [
249
+ "transfer_duration",
250
+ "post_duration",
251
+ "bytes_transferred",
252
+ "num_descriptors",
253
+ ],
254
+ ):
255
+ for list_item in transfer_stats_data[list_item_key]:
256
+ prom_obj[engine_idx].observe(list_item)
257
+ for counter_obj, counter_item_key in zip(
258
+ [
259
+ self.counter_nixl_num_failed_transfers,
260
+ self.counter_nixl_num_failed_notifications,
261
+ self.counter_nixl_num_kv_expired_reqs,
262
+ ],
263
+ ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"],
264
+ ):
265
+ for list_item in transfer_stats_data[counter_item_key]:
266
+ counter_obj[engine_idx].inc(list_item)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
new
+94
@@ -0,0 +1,94 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared constants, lazy imports and helpers for the NIXL connector."""
4
+
5
+import contextlib
6
+import os
7
+import sys
8
+from collections.abc import Iterator
9
+from typing import Any
10
+
11
+import zmq
12
+
13
+from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
+from vllm.utils.network_utils import make_zmq_socket
16
+
17
+logger = init_logger(__name__)
18
+
19
+
20
+# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
21
+try:
22
+ if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ:
23
+ # avoid a memory leak in UCX when using NIXL on some models
24
+ # see: https://github.com/vllm-project/vllm/issues/24264
25
+ if "nixl" in sys.modules or "rixl" in sys.modules:
26
+ logger.warning(
27
+ "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. "
28
+ "Please set it to '1024' manually."
29
+ )
30
+ else:
31
+ logger.info(
32
+ "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
33
+ "memory leak in UCX when using NIXL."
34
+ )
35
+ os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
36
+
37
+ if not current_platform.is_rocm():
38
+ from nixl._api import nixl_agent as NixlWrapper
39
+ from nixl._bindings import nixlXferTelemetry
40
+ else:
41
+ from rixl._api import nixl_agent as NixlWrapper
42
+ from rixl._bindings import nixlXferTelemetry
43
+
44
+ logger.info("NIXL is available")
45
+except ImportError:
46
+ logger.warning("NIXL is not available")
47
+ NixlWrapper = None
48
+ nixlXferTelemetry = None
49
+
50
+
51
+try:
52
+ if not current_platform.is_rocm():
53
+ from nixl._api import nixl_agent_config
54
+ else:
55
+ from rixl._api import nixl_agent_config
56
+except ImportError:
57
+ nixl_agent_config = None
58
+ logger.warning("NIXL agent config is not available")
59
+
60
+# Supported platforms and types of kv transfer buffer.
61
+# {device: tuple of supported kv buffer types}
62
+_NIXL_SUPPORTED_DEVICE = {
63
+ "cuda": (
64
+ "cuda",
65
+ "cpu",
66
+ ),
67
+ "tpu": ("cpu",),
68
+ "xpu": (
69
+ "cpu",
70
+ "xpu",
71
+ ),
72
+ "cpu": ("cpu",),
73
+}
74
+# support for oot platform by providing mapping in current_platform
75
+_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())
76
+
77
+
78
+# TODO: merge with vllm.utils.network_utils.zmq_socket_ctx
79
+@contextlib.contextmanager
80
+def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
81
+ """Context manager for a ZMQ socket"""
82
+
83
+ if socket_type not in (zmq.ROUTER, zmq.REQ):
84
+ raise ValueError(f"Unexpected socket type: {socket_type}")
85
+
86
+ ctx: zmq.Context | None = None
87
+ try:
88
+ ctx = zmq.Context() # type: ignore[attr-defined]
89
+ yield make_zmq_socket(
90
+ ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
91
+ )
92
+ finally:
93
+ if ctx is not None:
94
+ ctx.destroy(linger=0)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
renamed
+30
-1219
@@ -1,18 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import contextlib
4
-import copy
3
+"""Worker-side logic for the NIXL connector."""
4
+
5
import logging
6
import os
7
import queue
8
-import sys
8
import threading
9
import time
10
import uuid
11
from collections import defaultdict
13
-from collections.abc import Iterator
12
from concurrent.futures import Future, ThreadPoolExecutor
15
-from dataclasses import dataclass
13
from typing import TYPE_CHECKING, Any, cast
14
15
import msgspec
@@ -21,32 +18,36 @@ import torch
18
import zmq
19
20
from vllm import envs
24
-from vllm.config import VllmConfig
21
from vllm.distributed.kv_transfer.kv_connector.utils import (
22
BlockIds,
23
EngineId,
24
HeteroTPTransferConfig,
25
TpKVTopology,
30
- get_current_attn_backend,
26
get_current_attn_backends,
27
kv_postprocess_blksize_and_layout_on_receive,
28
kv_postprocess_blksize_on_receive,
29
kv_postprocess_layout_on_receive,
35
- yield_req_data,
30
)
37
-from vllm.distributed.kv_transfer.kv_connector.v1.base import (
38
- CopyBlocksOp,
39
- KVConnectorBase_V1,
40
- KVConnectorHandshakeMetadata,
41
- KVConnectorMetadata,
42
- KVConnectorRole,
43
- SupportsHMA,
31
+from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
32
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
33
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
34
+ GET_META_MSG,
35
+ NixlAgentMetadata,
36
+ NixlConnectorMetadata,
37
+ NixlHandshakePayload,
38
+ ReqId,
39
+ ReqMeta,
40
+ TransferHandle,
41
+ compute_nixl_compatibility_hash,
42
+)
43
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
44
+ NixlKVConnectorStats,
45
)
45
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
46
- KVConnectorPromMetrics,
47
- KVConnectorStats,
48
- PromMetric,
49
- PromMetricT,
46
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
47
+ _NIXL_SUPPORTED_DEVICE,
48
+ NixlWrapper,
49
+ nixl_agent_config,
50
+ zmq_ctx,
51
)
52
from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
53
MambaConvSplitInfo,
@@ -57,961 +58,34 @@ from vllm.distributed.parallel_state import (
58
get_tensor_model_parallel_rank,
59
get_tensor_model_parallel_world_size,
60
)
60
-from vllm.forward_context import ForwardContext
61
from vllm.logger import init_logger
62
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
63
from vllm.platforms import current_platform
64
-from vllm.utils.math_utils import cdiv
65
-from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
66
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
64
+from vllm.utils.network_utils import make_zmq_path
65
from vllm.v1.attention.backends.utils import get_kv_cache_layout
68
-from vllm.v1.core.sched.output import SchedulerOutput
66
from vllm.v1.kv_cache_interface import (
67
FullAttentionSpec,
68
MambaSpec,
72
- SlidingWindowSpec,
69
UniformTypeKVCacheSpecs,
70
)
75
-from vllm.v1.metrics.utils import create_metric_per_engine
71
from vllm.v1.worker.block_table import BlockTable
72
from vllm.v1.worker.utils import select_common_block_size
73
74
if TYPE_CHECKING:
80
- from vllm.v1.core.kv_cache_manager import KVCacheBlocks
75
+ from vllm.config import VllmConfig
76
from vllm.v1.kv_cache_interface import KVCacheConfig
82
- from vllm.v1.request import Request
83
-
84
-TransferHandle = int
85
-ReqId = str
86
-
87
-#
88
-# NIXL Connector Version
89
-#
90
-# Increment this version whenever there is an incompatible change to:
91
-# - NixlAgentMetadata schema
92
-# - kv_transfer_params schema or semantics
93
-# - NIXL transfer protocol or wire format
94
-# - KV cache memory layout or block organization
95
-# - Any other change that breaks P/D interoperability
96
-#
97
-# Version History:
98
-# 1: Initial version with compatibility checking
99
-# 2: Add remote_request_id to kv_transfer_params
100
-#
101
-NIXL_CONNECTOR_VERSION: int = 2
102
-
103
-GET_META_MSG = b"get_meta_msg"
77
78
logger = init_logger(__name__)
79
107
-# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
108
-try:
109
- if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ:
110
- # avoid a memory leak in UCX when using NIXL on some models
111
- # see: https://github.com/vllm-project/vllm/issues/24264
112
- if "nixl" in sys.modules or "rixl" in sys.modules:
113
- logger.warning(
114
- "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. "
115
- "Please set it to '1024' manually."
116
- )
117
- else:
118
- logger.info(
119
- "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
120
- "memory leak in UCX when using NIXL."
121
- )
122
- os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
123
-
124
- if not current_platform.is_rocm():
125
- from nixl._api import nixl_agent as NixlWrapper
126
- from nixl._bindings import nixlXferTelemetry
127
- else:
128
- from rixl._api import nixl_agent as NixlWrapper
129
- from rixl._bindings import nixlXferTelemetry
130
-
131
- logger.info("NIXL is available")
132
-except ImportError:
133
- logger.warning("NIXL is not available")
134
- NixlWrapper = None
135
- nixlXferTelemetry = None
136
-
137
-
138
-try:
139
- if not current_platform.is_rocm():
140
- from nixl._api import nixl_agent_config
141
- else:
142
- from rixl._api import nixl_agent_config
143
-except ImportError:
144
- nixl_agent_config = None
145
- logger.warning("NIXL agent config is not available")
146
-
147
-# Supported platforms and types of kv transfer buffer.
148
-# {device: tuple of supported kv buffer types}
149
-_NIXL_SUPPORTED_DEVICE = {
150
- "cuda": (
151
- "cuda",
152
- "cpu",
153
- ),
154
- "tpu": ("cpu",),
155
- "xpu": (
156
- "cpu",
157
- "xpu",
158
- ),
159
- "cpu": ("cpu",),
160
-}
161
-# support for oot platform by providing mapping in current_platform
162
-_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())
163
-
164
-
165
-@dataclass
166
-class NixlAgentMetadata:
167
- engine_id: str
168
- agent_metadata: bytes
169
- kv_caches_base_addr: list[int]
170
- device_id: int
171
- num_blocks: int
172
- block_lens: list[int]
173
- kv_cache_layout: str
174
- block_size: int
175
- ssm_sizes: tuple[int, int]
176
- attn_backend_name: str
177
-
178
-
179
-@dataclass
180
-class NixlHandshakePayload(KVConnectorHandshakeMetadata):
181
- """
182
- Wrapper for NIXL handshake sent over the wire.
183
-
184
- Enables two-phase decoding for graceful compatibility checking:
185
- 1. Decode NixlHandshakePayload to get compatibility_hash
186
- 2. Compute local hash and compare
187
- 3. Only if hashes match, decode agent_metadata_bytes
188
-
189
- This prevents decoder errors when NixlAgentMetadata schema is
190
- incompatible, allowing graceful failure with clear error message.
191
- """
192
-
193
- compatibility_hash: str
194
- agent_metadata_bytes: bytes # NixlAgentMetadata encoded
195
-
196
-
197
-def compute_nixl_compatibility_hash(
198
- vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
199
-) -> str:
200
- """
201
- Compute compatibility hash for NIXL KV transfer.
202
-
203
- Hash only the factors that affect whether two NIXL instances can
204
- successfully transfer KV cache data.
205
-
206
- Factors included:
207
- - vLLM version and NIXL connector version
208
- - Model architecture (name, dtype, KV heads, layers)
209
- - KV cache format (dtype, sliding window)
210
- - Attention backend
211
-
212
- Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout
213
- are validated at runtime in _validate_remote_agent_handshake and are not
214
- included in this hash to support heterogeneous deployments.
215
-
216
- Note - the set of factors are likely to evolve significantly over
217
- time to be more or less permissive.
218
-
219
- Returns:
220
- SHA-256 hex digest
221
- """
222
- from vllm import __version__ as vllm_version
223
- from vllm.config.utils import hash_factors
224
-
225
- model_config = vllm_config.model_config
226
- cache_config = vllm_config.cache_config
227
- is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
228
-
229
- factors = {
230
- # Version compatibility
231
- "vllm_version": vllm_version,
232
- "nixl_connector_version": NIXL_CONNECTOR_VERSION,
233
- # Model architecture - affects KV cache shape
234
- "model": model_config.model,
235
- "dtype": str(model_config.dtype),
236
- "num_kv_heads": model_config.get_total_num_kv_heads(),
237
- "head_size": model_config.get_head_size(),
238
- "num_hidden_layers": model_config.get_total_num_hidden_layers(),
239
- # Attention backend and KV cache dtype affect memory layout
240
- "attn_backend_name": attn_backend_name,
241
- "cache_dtype": str(cache_config.cache_dtype),
242
- "cross_layers_blocks": cross_layers_blocks,
243
- "is_hma_enabled": is_hma_enabled,
244
- }
245
-
246
- compat_hash = hash_factors(factors)
247
- logger.debug(
248
- "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, "
249
- "cache_dtype=%s, attn_backend=%s)",
250
- compat_hash,
251
- factors["model"],
252
- factors["dtype"],
253
- factors["num_kv_heads"],
254
- factors["cache_dtype"],
255
- attn_backend_name,
256
- )
257
- return compat_hash
258
-
259
-
260
-@dataclass
261
-class RemoteMeta:
262
- block_ids: BlockIds
263
- host: str
264
- port: int
265
- engine_id: str
266
- request_id: str
267
-
268
-
269
-@dataclass
270
-class ReqMeta:
271
- local_block_ids: BlockIds
272
- # To be used when logical block size does not match the kernel block size
273
- local_physical_block_ids: BlockIds
274
- tp_size: int
275
- remote: RemoteMeta | None = None
276
-
277
-
278
-class NixlConnectorMetadata(KVConnectorMetadata):
279
- def __init__(self):
280
- self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
281
- self.reqs_to_save: dict[ReqId, ReqMeta] = {}
282
- self.reqs_to_send: dict[ReqId, float] = {}
283
- self.reqs_in_batch: set[ReqId] = set()
284
- self.reqs_not_processed: set[ReqId] = set()
285
-
286
- def _add_new_req(
287
- self,
288
- local_block_ids: BlockIds,
289
- kv_transfer_params: dict[str, Any],
290
- ) -> ReqMeta:
291
- return ReqMeta(
292
- local_block_ids=local_block_ids,
293
- local_physical_block_ids=local_block_ids,
294
- # P workers don't need to receive tp_size from proxy here.
295
- tp_size=kv_transfer_params.get("tp_size", 1),
296
- )
297
-
298
- def add_new_req_to_save(
299
- self,
300
- request_id: ReqId,
301
- local_block_ids: BlockIds,
302
- kv_transfer_params: dict[str, Any],
303
- ):
304
- self.reqs_to_save[request_id] = self._add_new_req(
305
- local_block_ids, kv_transfer_params
306
- )
307
-
308
- def add_new_req_to_recv(
309
- self,
310
- request_id: ReqId,
311
- local_block_ids: BlockIds,
312
- kv_transfer_params: dict[str, Any],
313
- ):
314
- req = self._add_new_req(local_block_ids, kv_transfer_params)
315
- req.remote = RemoteMeta(
316
- block_ids=kv_transfer_params["remote_block_ids"],
317
- engine_id=kv_transfer_params["remote_engine_id"],
318
- request_id=kv_transfer_params["remote_request_id"],
319
- host=kv_transfer_params["remote_host"],
320
- port=kv_transfer_params["remote_port"],
321
- )
322
- self.reqs_to_recv[request_id] = req
323
-
324
-
325
-class NixlConnector(KVConnectorBase_V1, SupportsHMA):
326
- @property
327
- def prefer_cross_layer_blocks(self) -> bool:
328
- if any(
329
- [
330
- isinstance(group.kv_cache_spec, MambaSpec)
331
- for group in self.kv_cache_config.kv_cache_groups
332
- ]
333
- ):
334
- # Hybrid SSM models do not yet support cross-layer layout
335
- return False
336
-
337
- backend = get_current_attn_backend(self._vllm_config)
338
- if backend.get_name() not in (
339
- "FLASH_ATTN",
340
- "FLASHINFER",
341
- "TRITON_ATTN",
342
- ):
343
- return False
344
-
345
- # For now there is no benefit to run cross layers when backend
346
- # does not support on HND
347
- if get_kv_cache_layout() != "HND":
348
- return False
349
-
350
- extra_config = self.kv_transfer_config.kv_connector_extra_config
351
- return (
352
- str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
353
- == "true"
354
- )
355
-
356
- def __init__(
357
- self,
358
- vllm_config: VllmConfig,
359
- role: KVConnectorRole,
360
- kv_cache_config: "KVCacheConfig",
361
- ):
362
- super().__init__(vllm_config, role, kv_cache_config)
363
- assert vllm_config.kv_transfer_config is not None
364
- assert vllm_config.kv_transfer_config.engine_id is not None
365
- self.kv_cache_config = kv_cache_config
366
- self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
367
- self.kv_transfer_config = vllm_config.kv_transfer_config
368
- if role == KVConnectorRole.SCHEDULER:
369
- self.connector_scheduler: NixlConnectorScheduler | None = (
370
- NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
371
- )
372
- self.connector_worker: NixlConnectorWorker | None = None
373
- elif role == KVConnectorRole.WORKER:
374
- self.connector_scheduler = None
375
- self.connector_worker = NixlConnectorWorker(
376
- vllm_config, self.engine_id, kv_cache_config
377
- )
378
-
379
- ############################################################
380
- # Class Methods
381
- ############################################################
382
- @classmethod
383
- def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
384
- if vllm_config.model_config is None:
385
- logger.warning_once(
386
- "Unable to detect current VLLM config. "
387
- "Fallback to default kv cache layout."
388
- )
389
- return None
390
- use_mla = vllm_config.model_config.use_mla
391
- if use_mla:
392
- # return None when we have mla
393
- # as the layout should not matter in that case,
394
- # which fallback to the default behavior.
395
- return None
396
- logger.info_once(
397
- "NixlConnector setting KV cache layout to HND for better xfer performance."
398
- )
399
- return "HND"
400
-
401
- ############################################################
402
- # Scheduler Side Methods
403
- ############################################################
404
-
405
- def get_num_new_matched_tokens(
406
- self, request: "Request", num_computed_tokens: int
407
- ) -> tuple[int | None, bool]:
408
- assert self.connector_scheduler is not None
409
- return self.connector_scheduler.get_num_new_matched_tokens(
410
- request, num_computed_tokens
411
- )
412
-
413
- def update_state_after_alloc(
414
- self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
415
- ):
416
- assert self.connector_scheduler is not None
417
- return self.connector_scheduler.update_state_after_alloc(
418
- request, blocks, num_external_tokens
419
- )
420
-
421
- def build_connector_meta(
422
- self,
423
- scheduler_output: SchedulerOutput,
424
- ) -> KVConnectorMetadata:
425
- assert self.connector_scheduler is not None
426
- return self.connector_scheduler.build_connector_meta(scheduler_output)
427
-
428
- def request_finished(
429
- self,
430
- request: "Request",
431
- block_ids: list[int],
432
- ) -> tuple[bool, dict[str, Any] | None]:
433
- assert self.connector_scheduler is not None
434
- return self.connector_scheduler.request_finished(request, (block_ids,))
435
-
436
- def request_finished_all_groups(
437
- self,
438
- request: "Request",
439
- block_ids: tuple[list[int], ...],
440
- ) -> tuple[bool, dict[str, Any] | None]:
441
- assert self.connector_scheduler is not None
442
- return self.connector_scheduler.request_finished(request, block_ids)
443
-
444
- def set_xfer_handshake_metadata(
445
- self, metadata: dict[int, KVConnectorHandshakeMetadata]
446
- ) -> None:
447
- """
448
- Set the KV connector handshake metadata for this connector.
449
-
450
- Args:
451
- metadata (dict): the handshake metadata to set.
452
- """
453
- assert self.connector_scheduler is not None
454
- self.connector_scheduler.set_xfer_handshake_metadata(metadata)
455
-
456
- ############################################################
457
- # Worker Side Methods
458
- ############################################################
459
- def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
460
- assert self.connector_worker is not None
461
- self.connector_worker.register_kv_caches(kv_caches)
462
-
463
- def register_cross_layers_kv_cache(
464
- self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
465
- ):
466
- assert self.connector_worker is not None
467
- self.connector_worker.register_cross_layers_kv_caches(kv_cache)
468
-
469
- def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
470
- assert self.connector_worker is not None
471
- self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
472
-
473
- def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
474
- """Get the finished recving and sending requests."""
475
- assert self.connector_worker is not None
476
- return self.connector_worker.get_finished()
477
-
478
- def get_block_ids_with_load_errors(self) -> set[int]:
479
- """Get block IDs that failed to load via NIXL."""
480
- assert self.connector_worker is not None
481
- return self.connector_worker.get_block_ids_with_load_errors()
482
-
483
- def get_kv_connector_stats(self) -> KVConnectorStats | None:
484
- if self.connector_worker is None:
485
- return None
486
- return self.connector_worker.get_kv_connector_stats()
487
-
488
- @classmethod
489
- def build_kv_connector_stats(
490
- cls, data: dict[str, Any] | None = None
491
- ) -> KVConnectorStats | None:
492
- return (
493
- NixlKVConnectorStats(data=data)
494
- if data is not None
495
- else NixlKVConnectorStats()
496
- )
497
-
498
- @classmethod
499
- def build_prom_metrics(
500
- cls,
501
- vllm_config: VllmConfig,
502
- metric_types: dict[type[PromMetric], type[PromMetricT]],
503
- labelnames: list[str],
504
- per_engine_labelvalues: dict[int, list[object]],
505
- ) -> KVConnectorPromMetrics:
506
- return NixlPromMetrics(
507
- vllm_config, metric_types, labelnames, per_engine_labelvalues
508
- )
509
-
510
- def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
511
- assert self.connector_worker is not None
512
- assert isinstance(self._connector_metadata, NixlConnectorMetadata)
513
- self.connector_worker.start_load_kv(self._connector_metadata)
514
-
515
- def wait_for_layer_load(self, layer_name: str) -> None:
516
- """NixlConnector does not do layerwise saving."""
517
- pass
518
-
519
- def save_kv_layer(
520
- self,
521
- layer_name: str,
522
- kv_layer: torch.Tensor,
523
- attn_metadata: AttentionMetadata,
524
- **kwargs,
525
- ) -> None:
526
- """NixlConnector does not save explicitly."""
527
- pass
528
-
529
- def wait_for_save(self):
530
- assert self.connector_worker is not None
531
- assert isinstance(self._connector_metadata, NixlConnectorMetadata)
532
- if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
533
- self.connector_worker.save_kv_to_host(self._connector_metadata)
534
-
535
- def shutdown(self):
536
- if self.connector_worker is not None:
537
- self.connector_worker.shutdown()
538
- if self.connector_scheduler is not None:
539
- self.connector_scheduler.shutdown()
540
-
541
- def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
542
- """
543
- Get the KVConnector handshake metadata for this connector.
544
- This metadata is used for out-of-band connector handshake
545
- between P/D workers.
546
-
547
- Returns:
548
- KVConnectorHandshakeMetadata: the handshake metadata.
549
- None if no handshake metadata is available.
550
- """
551
- assert self.connector_worker is not None
552
- return self.connector_worker.xfer_handshake_metadata
553
-
554
-
555
-class NixlConnectorScheduler:
556
- """Implementation of Scheduler side methods"""
557
-
558
- def __init__(
559
- self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig"
560
- ):
561
- self.vllm_config = vllm_config
562
- self.block_size = vllm_config.cache_config.block_size
563
- self.engine_id: EngineId = engine_id
564
- self.kv_cache_config = kv_cache_config
565
- self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
566
- self.side_channel_port = (
567
- envs.VLLM_NIXL_SIDE_CHANNEL_PORT
568
- + vllm_config.parallel_config.data_parallel_index
569
- )
570
- assert vllm_config.kv_transfer_config is not None
571
- if current_platform.device_type == "cpu":
572
- self.use_host_buffer = False
573
- else:
574
- self.use_host_buffer = (
575
- vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
576
- )
577
- self._is_hma_required = (
578
- not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
579
- # Also handle unlikely SW-only model case instead of checking num_groups>1.
580
- and any(
581
- not isinstance(g.kv_cache_spec, FullAttentionSpec)
582
- for g in kv_cache_config.kv_cache_groups
583
- )
584
- )
585
- self._has_mamba = any(
586
- isinstance(g.kv_cache_spec, MambaSpec)
587
- for g in kv_cache_config.kv_cache_groups
588
- )
589
-
590
- logger.info("Initializing NIXL Scheduler %s", engine_id)
591
- if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
592
- logger.info("Hybrid Memory Allocator is enabled with NIXL")
593
-
594
- # Background thread for handling new handshake requests.
595
- self._nixl_handshake_listener_t: threading.Thread | None = None
596
- self._stop_event = threading.Event()
597
-
598
- # Requests that need to start recv/send.
599
- # New requests are added by update_state_after_alloc in
600
- # the scheduler. Used to make metadata passed to Worker.
601
- self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
602
- self._reqs_need_save: dict[ReqId, Request] = {}
603
- # Reqs to send and their expiration time
604
- self._reqs_need_send: dict[ReqId, float] = {}
605
- self._reqs_in_batch: set[ReqId] = set()
606
- # Reqs to remove from processed set because they're not to send after
607
- # remote prefill or aborted.
608
- self._reqs_not_processed: set[ReqId] = set()
609
-
610
- # Gather Sliding Window sizes for each kv cache group (if any) in number of
611
- # blocks per KV cache group. This is used to clip the local attention window.
612
- sw_sizes_tokens: list[tuple[int, int]] = [
613
- (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
614
- if isinstance(g.kv_cache_spec, SlidingWindowSpec)
615
- else (0, self.block_size)
616
- for g in kv_cache_config.kv_cache_groups
617
- ]
618
- # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
619
- # account for boundary overlap eg window isn't fully aligned with blocks.
620
- self.blocks_per_sw = [
621
- cdiv(n_tokens, block_size) + 1 if n_tokens else 0
622
- for n_tokens, block_size in sw_sizes_tokens
623
- ]
624
-
625
- def shutdown(self):
626
- self._stop_event.set()
627
- if self._nixl_handshake_listener_t is not None:
628
- self._nixl_handshake_listener_t.join()
629
- self._nixl_handshake_listener_t = None
630
-
631
- def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
632
- """
633
- Clip the number of blocks to the sliding window size for each kv cache group
634
- that employs SWA.
635
- This is necessary because the KV Cache manager initially allocates blocks for
636
- the entire sequence length, and successively cleans up blocks that are outside
637
- the window prior to the `request_finished_all_groups` hook.
638
- """
639
- if len(block_ids) == 0 or not self._is_hma_required:
640
- # No blocks to clip eg Full prefix cache hit or not a hybrid model.
641
- return block_ids
642
- # NOTE (NickLucche) This logic is currently handled at the connector level
643
- # because offloading connectors might want to receive the whole sequence even
644
- # for SWA groups. We will abstract this logic once the interface is more stable
645
- assert len(block_ids) == len(self.blocks_per_sw), (
646
- "Number of KV cache groups must match"
647
- )
648
- # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
649
- return tuple(
650
- [
651
- blocks[-self.blocks_per_sw[i] :]
652
- if self.blocks_per_sw[i] > 0
653
- else blocks
654
- for i, blocks in enumerate(block_ids)
655
- ]
656
- )
657
-
658
- def set_xfer_handshake_metadata(
659
- self, metadata: dict[int, KVConnectorHandshakeMetadata]
660
- ) -> None:
661
- """
662
- Set the KV connector handshake metadata for this connector.
663
-
664
- Args:
665
- metadata (dict): the handshake metadata to set.
666
- """
667
- encoded_data: dict[int, bytes] = {}
668
- encoder = msgspec.msgpack.Encoder()
669
- for tp_rank, rank_metadata in metadata.items():
670
- if not isinstance(rank_metadata, NixlHandshakePayload):
671
- raise ValueError(
672
- "NixlConnectorScheduler expects NixlHandshakePayload for "
673
- "handshake metadata."
674
- )
675
- encoded_data[tp_rank] = encoder.encode(rank_metadata)
676
- logger.debug(
677
- "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",
678
- tp_rank,
679
- str(len(encoded_data[tp_rank])),
680
- )
681
-
682
- # Only start the listener when we have metadata to serve.
683
- if self._nixl_handshake_listener_t is None:
684
- ready_event = threading.Event()
685
- self._nixl_handshake_listener_t = threading.Thread(
686
- target=self._nixl_handshake_listener,
687
- args=(
688
- encoded_data,
689
- ready_event,
690
- self._stop_event,
691
- self.side_channel_port,
692
- ),
693
- daemon=True,
694
- name="nixl_handshake_listener",
695
- )
696
- self._nixl_handshake_listener_t.start()
697
- ready_event.wait() # Wait for listener ZMQ socket to be ready.
698
-
699
- @staticmethod
700
- def _nixl_handshake_listener(
701
- encoded_data: dict[int, Any],
702
- ready_event: threading.Event,
703
- stop_event: threading.Event,
704
- port: int,
705
- ):
706
- """Background thread for getting new NIXL handshakes."""
707
- # NOTE(rob): this is a simple implementation. We will move
708
- # to a better approach via HTTP endpoint soon.
709
-
710
- # Listen for new requests for metadata.
711
- host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
712
- path = make_zmq_path("tcp", host, port)
713
- logger.debug("Starting listening on path: %s", path)
714
- with zmq_ctx(zmq.ROUTER, path) as sock:
715
- sock.setsockopt(zmq.RCVTIMEO, 1000)
716
- ready_event.set()
717
- while True:
718
- try:
719
- identity, _, msg = sock.recv_multipart()
720
- except zmq.Again:
721
- if stop_event.is_set():
722
- break
723
- continue
724
- # Decode the message which contains (GET_META_MSG, rank)
725
- msg, target_tp_rank = msgspec.msgpack.decode(msg)
726
- logger.debug(
727
- "Received message for tp rank %s",
728
- target_tp_rank,
729
- )
730
- if msg != GET_META_MSG:
731
- logger.warning("Connection listener got unexpected message %s", msg)
732
- sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
733
-
734
- def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
735
- """D-side only. Returns N-1 for Mamba models since the decoder
736
- always recomputes the last token and must start from h(N-1)."""
737
- if self._has_mamba and num_prompt_tokens > 1:
738
- return num_prompt_tokens - 1
739
- return num_prompt_tokens
740
-
741
- def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
742
- """P-side only: drop the last prompt token so the prefiller computes
743
- h(N-1) instead of h(N). The decoder recomputes the last token to
744
- derive h(N) correctly.
745
-
746
- Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
747
- request is preempted and rescheduled."""
748
- params = request.kv_transfer_params
749
- if (
750
- params is not None
751
- # Guard against repeated truncation after preemption/reschedule.
752
- and not params.get("_p_side_truncated")
753
- and request.num_prompt_tokens > 1
754
- ):
755
- if request.prompt_token_ids is not None:
756
- request.prompt_token_ids.pop()
757
- elif request.prompt_embeds is not None:
758
- request.prompt_embeds = request.prompt_embeds[:-1]
759
- else:
760
- return
761
-
762
- request._all_token_ids.pop()
763
- request.num_prompt_tokens -= 1
764
- request.max_tokens = 1
765
- params["_p_side_truncated"] = True
766
-
767
- def get_num_new_matched_tokens(
768
- self, request: "Request", num_computed_tokens: int
769
- ) -> tuple[int, bool]:
770
- """
771
- For remote prefill, pull all prompt blocks from remote
772
- asynchronously relative to engine execution.
773
-
774
- Args:
775
- request (Request): the request object.
776
- num_computed_tokens (int): the number of locally
777
- computed tokens for this request
778
- Returns:
779
- * the number of tokens that can be loaded from the
780
- external KV cache beyond what is already computed.
781
- * true if the external KV cache tokens will be loaded
782
- asynchronously (between scheduler steps).
783
- """
784
-
785
- params = request.kv_transfer_params
786
- logger.debug(
787
- "NIXLConnector get_num_new_matched_tokens: "
788
- "num_computed_tokens=%s, kv_transfer_params=%s",
789
- num_computed_tokens,
790
- params,
791
- )
792
-
793
- if params is not None and params.get("do_remote_prefill"):
794
- # Remote prefill: get all prompt blocks from remote.
795
- token_ids = request.prompt_token_ids or []
796
- actual = self._mamba_prefill_token_count(len(token_ids))
797
- count = actual - num_computed_tokens
798
- if count > 0:
799
- return count, True
800
-
801
- if params is not None and params.get("do_remote_decode") and self._has_mamba:
802
- self._truncate_mamba_request_for_prefill(request)
803
-
804
- # No remote prefill for this request.
805
- return 0, False
806
-
807
- def update_state_after_alloc(
808
- self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
809
- ):
810
- params = request.kv_transfer_params
811
- logger.debug(
812
- "NIXLConnector update_state_after_alloc: "
813
- "num_external_tokens=%s, kv_transfer_params=%s",
814
- num_external_tokens,
815
- params,
816
- )
817
-
818
- if not params:
819
- return
820
-
821
- if params.get("do_remote_decode"):
822
- self._reqs_in_batch.add(request.request_id)
823
- if self.use_host_buffer and params.get("do_remote_decode"):
824
- # NOTE: when accelerator is not directly supported by Nixl,
825
- # prefilled blocks need to be saved to host memory before transfer.
826
- self._reqs_need_save[request.request_id] = request
827
- elif params.get("do_remote_prefill"):
828
- if params.get("remote_block_ids"):
829
- if all(
830
- p in params
831
- for p in (
832
- "remote_engine_id",
833
- "remote_request_id",
834
- "remote_host",
835
- "remote_port",
836
- )
837
- ):
838
- # If remote_blocks and num_external_tokens = 0, we have
839
- # a full prefix cache hit on the D worker. We need to call
840
- # send_notif in _read_blocks to free the memory on the P.
841
-
842
- unhashed_local_block_ids: BlockIds = (
843
- blocks.get_unhashed_block_ids_all_groups()
844
- if num_external_tokens > 0
845
- else ()
846
- )
847
- local_block_ids = self.get_sw_clipped_blocks(
848
- unhashed_local_block_ids
849
- )
850
-
851
- # Get unhashed blocks to pull from remote. Mind that a full prefix
852
- # cache hit is indicated with an empty list.
853
- self._reqs_need_recv[request.request_id] = (
854
- request,
855
- local_block_ids,
856
- )
857
-
858
- else:
859
- logger.warning(
860
- "Got invalid KVTransferParams: %s. This "
861
- "request will not utilize KVTransfer",
862
- params,
863
- )
864
- else:
865
- assert num_external_tokens == 0
866
- # Only trigger 1 KV transfer per request.
867
- params["do_remote_prefill"] = False
868
-
869
- def _build_save_meta(
870
- self,
871
- meta: NixlConnectorMetadata,
872
- scheduler_output: SchedulerOutput,
873
- ) -> None:
874
- # only called when use_host_buffer is True to build the save metadata
875
-
876
- # NOTE: For the prefill side, there might be a chance that an early added
877
- # request is a chunked prefill, so we need to check if new blocks are added
878
- for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
879
- req_to_save = self._reqs_need_save.get(req_id)
880
- if req_to_save is None or new_block_id_groups is None:
881
- continue
882
- req = req_to_save
883
-
884
- assert req.kv_transfer_params is not None
885
- clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
886
- meta.add_new_req_to_save(
887
- request_id=req_id,
888
- local_block_ids=clipped_block_id_groups,
889
- kv_transfer_params=req.kv_transfer_params,
890
- )
891
- assert scheduler_output.num_scheduled_tokens is not None
892
- num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
893
- is_partial = (
894
- req.num_computed_tokens + num_scheduled_tokens
895
- ) < req.num_prompt_tokens
896
- if not is_partial:
897
- # For non-partial prefills, once new req_meta is scheduled, it
898
- # can be removed from _reqs_need_save.
899
- # For partial prefill case, we will retain the request in
900
- # _reqs_need_save until all blocks are scheduled with req_meta.
901
- # Therefore, only pop if `not is_partial`.
902
- self._reqs_need_save.pop(req_id)
903
-
904
- def build_connector_meta(
905
- self,
906
- scheduler_output: SchedulerOutput,
907
- ) -> KVConnectorMetadata:
908
- meta = NixlConnectorMetadata()
909
-
910
- # Loop through scheduled reqs and convert to ReqMeta.
911
- for req_id, (req, block_ids) in self._reqs_need_recv.items():
912
- assert req.kv_transfer_params is not None
913
- meta.add_new_req_to_recv(
914
- request_id=req_id,
915
- local_block_ids=block_ids,
916
- kv_transfer_params=req.kv_transfer_params,
917
- )
918
-
919
- if self.use_host_buffer:
920
- self._build_save_meta(meta, scheduler_output)
921
-
922
- meta.reqs_to_send = self._reqs_need_send
923
- meta.reqs_in_batch = self._reqs_in_batch
924
- meta.reqs_not_processed = self._reqs_not_processed
925
-
926
- # Clear the list once workers start the transfers
927
- self._reqs_need_recv.clear()
928
- self._reqs_in_batch = set()
929
- self._reqs_not_processed = set()
930
- self._reqs_need_send = {}
931
-
932
- return meta
933
-
934
- def request_finished(
935
- self,
936
- request: "Request",
937
- block_ids: BlockIds,
938
- ) -> tuple[bool, dict[str, Any] | None]:
939
- """
940
- Once a request is finished, determine whether request blocks
941
- should be freed now or will be sent asynchronously and freed later.
942
- """
943
- from vllm.v1.request import RequestStatus
944
-
945
- params = request.kv_transfer_params
946
- logger.debug(
947
- "NIXLConnector request_finished(%s), request_status=%s, "
948
- "kv_transfer_params=%s",
949
- request.request_id,
950
- request.status,
951
- params,
952
- )
953
- if not params:
954
- return False, None
955
-
956
- if params.get("do_remote_prefill"):
957
- # If do_remote_prefill is still True when the request is finished,
958
- # update_state_after_alloc must not have been called (the request
959
- # must have been aborted before it was scheduled).
960
- # To avoid stranding the prefill blocks in the prefill instance,
961
- # we must add empty block_ids to _reqs_need_recv so that our
962
- # worker side will notify and free blocks in the prefill instance.
963
- self._reqs_need_recv[request.request_id] = (request, [])
964
- params["do_remote_prefill"] = False
965
- return False, None
966
-
967
- if not params.get("do_remote_decode"):
968
- return False, None
969
- if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
970
- # Also include the case of a P/D Prefill request with immediate
971
- # block free (eg abort). Stop tracking this request.
972
- self._reqs_not_processed.add(request.request_id)
973
- # Clear _reqs_need_save if a request is aborted as partial prefill.
974
- self._reqs_need_save.pop(request.request_id, None)
975
- return False, None
976
-
977
- # TODO: check whether block_ids actually ever be 0. If not we could
978
- # remove the conditional below
979
- delay_free_blocks = any(len(group) > 0 for group in block_ids)
980
-
981
- if delay_free_blocks:
982
- # Prefill request on remote. It will be read from D upon completion
983
- logger.debug(
984
- "NIXLConnector request_finished(%s) waiting for %d seconds "
985
- "for remote decode to fetch blocks",
986
- request.request_id,
987
- envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT,
988
- )
989
- self._reqs_need_send[request.request_id] = (
990
- time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
991
- )
992
- # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
993
- # trimming down after allocating for the whole sequence length. Empty
994
- # blocks are always at the start of the list.
995
- # Here we "unpad" blocks to send the actual remote blocks to be read.
996
- block_ids = self.get_sw_clipped_blocks(block_ids)
997
-
998
- return delay_free_blocks, dict(
999
- do_remote_prefill=True,
1000
- do_remote_decode=False,
1001
- remote_block_ids=block_ids,
1002
- remote_engine_id=self.engine_id,
1003
- remote_request_id=request.request_id,
1004
- remote_host=self.side_channel_host,
1005
- remote_port=self.side_channel_port,
1006
- tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
1007
- )
1008
-
80
81
class NixlConnectorWorker:
82
"""Implementation of Worker side methods"""
83
84
def __init__(
1014
- self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig"
85
+ self,
86
+ vllm_config: "VllmConfig",
87
+ engine_id: str,
88
+ kv_cache_config: "KVCacheConfig",
89
):
90
if NixlWrapper is None:
91
logger.error("NIXL is not available")
@@ -3200,9 +2274,9 @@ class NixlConnectorWorker:
2274
the their size differs.
2275
Reference diagram:
2276
KVCacheTensor (Shared)
3203
- / \
3204
- / \
3205
- / \
2277
+ / \\
2278
+ / \\
2279
+ / \\
2280
Attention (FlashInfer) View Mamba View
2281
| |
2282
| |
@@ -3283,266 +2357,3 @@ class NixlConnectorWorker:
2357
for desc in self._registered_descs:
2358
self.nixl_wrapper.deregister_memory(desc)
2359
self._registered_descs.clear()
3286
-
3287
-
3288
-@contextlib.contextmanager
3289
-def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
3290
- """Context manager for a ZMQ socket"""
3291
-
3292
- if socket_type not in (zmq.ROUTER, zmq.REQ):
3293
- raise ValueError(f"Unexpected socket type: {socket_type}")
3294
-
3295
- ctx: zmq.Context | None = None
3296
- try:
3297
- ctx = zmq.Context() # type: ignore[attr-defined]
3298
- yield make_zmq_socket(
3299
- ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
3300
- )
3301
- finally:
3302
- if ctx is not None:
3303
- ctx.destroy(linger=0)
3304
-
3305
-
3306
-@dataclass
3307
-class NixlKVConnectorStats(KVConnectorStats):
3308
- """Container for transfer performance metrics"""
3309
-
3310
- def __post_init__(self):
3311
- if not self.data:
3312
- # Empty container init, no data is passed in.
3313
- self.reset()
3314
-
3315
- def reset(self):
3316
- # Must be serializable
3317
- self.data: dict[str, list[float | int]] = {
3318
- "transfer_duration": [],
3319
- "post_duration": [],
3320
- "bytes_transferred": [],
3321
- "num_descriptors": [],
3322
- "num_failed_transfers": [],
3323
- "num_failed_notifications": [],
3324
- "num_kv_expired_reqs": [],
3325
- }
3326
-
3327
- def record_transfer(self, res: nixlXferTelemetry):
3328
- # Keep metrics units consistent with rest of the code: time us->s
3329
- self.data["transfer_duration"].append(res.xferDuration / 1e6)
3330
- self.data["post_duration"].append(res.postDuration / 1e6)
3331
- self.data["bytes_transferred"].append(res.totalBytes)
3332
- self.data["num_descriptors"].append(res.descCount)
3333
-
3334
- def record_failed_transfer(self):
3335
- """Record a failed NIXL transfer operation."""
3336
- self.data["num_failed_transfers"].append(1)
3337
-
3338
- def record_failed_notification(self):
3339
- """Record a failed NIXL notification (send_notif)."""
3340
- self.data["num_failed_notifications"].append(1)
3341
-
3342
- def record_kv_expired_req(self):
3343
- """Record a request that had its KV blocks expire."""
3344
- self.data["num_kv_expired_reqs"].append(1)
3345
-
3346
- def clone_and_reset(self) -> "NixlKVConnectorStats":
3347
- old = copy.copy(self)
3348
- self.reset()
3349
- return old
3350
-
3351
- def is_empty(self) -> bool:
3352
- # Do not discard metrics update that are entirely failures related.
3353
- return (
3354
- self.num_successful_transfers == 0
3355
- and len(self.data["num_failed_transfers"]) == 0
3356
- and len(self.data["num_failed_notifications"]) == 0
3357
- and len(self.data["num_kv_expired_reqs"]) == 0
3358
- )
3359
-
3360
- def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
3361
- if not other.is_empty():
3362
- for k, v in other.data.items():
3363
- accumulator = self.data[k]
3364
- assert isinstance(accumulator, list)
3365
- accumulator.extend(v)
3366
- return self
3367
-
3368
- def reduce(self) -> dict[str, int | float]:
3369
- # Compute compact representative stats suitable for CLI logging
3370
- if self.num_successful_transfers == 0:
3371
- # CLI logging only reports successful transfers stats. If all requests in
3372
- # the interval were unsuccessful, Prom will report failures stats instead.
3373
- return {
3374
- "Num successful transfers": 0,
3375
- "Avg xfer time (ms)": 0,
3376
- "P90 xfer time (ms)": 0,
3377
- "Avg post time (ms)": 0,
3378
- "P90 post time (ms)": 0,
3379
- "Avg MB per transfer": 0,
3380
- "Throughput (MB/s)": 0,
3381
- "Avg number of descriptors": 0,
3382
- }
3383
-
3384
- xfer_time = np.asarray(self.data["transfer_duration"])
3385
- post_time = np.asarray(self.data["post_duration"])
3386
- # Convert to MB for CLI logging.
3387
- mb = np.asarray(self.data["bytes_transferred"]) / 2**20
3388
- descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32)
3389
- n = len(descs)
3390
- assert n == self.num_successful_transfers
3391
-
3392
- total_mb = mb.sum()
3393
- avg_mb = total_mb / n
3394
-
3395
- total_time_seconds = xfer_time.sum()
3396
- throughput_mb_s = total_mb / total_time_seconds
3397
-
3398
- return {
3399
- "Num successful transfers": n,
3400
- "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
3401
- "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
3402
- "Avg post time (ms)": round(post_time.mean() * 1e3, 3),
3403
- "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),
3404
- "Avg MB per transfer": round(avg_mb, 3),
3405
- "Throughput (MB/s)": round(throughput_mb_s, 3),
3406
- "Avg number of descriptors": round(descs.mean(), 1),
3407
- }
3408
-
3409
- @property
3410
- def num_successful_transfers(self) -> int:
3411
- return len(self.data["transfer_duration"])
3412
-
3413
-
3414
-class NixlPromMetrics(KVConnectorPromMetrics):
3415
- def __init__(
3416
- self,
3417
- vllm_config: VllmConfig,
3418
- metric_types: dict[type[PromMetric], type[PromMetricT]],
3419
- labelnames: list[str],
3420
- per_engine_labelvalues: dict[int, list[object]],
3421
- ):
3422
- super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
3423
-
3424
- buckets = [
3425
- 0.001,
3426
- 0.005,
3427
- 0.01,
3428
- 0.025,
3429
- 0.05,
3430
- 0.075,
3431
- 0.1,
3432
- 0.2,
3433
- 0.3,
3434
- 0.5,
3435
- 0.75,
3436
- 1.0,
3437
- 5.0,
3438
- ]
3439
- nixl_histogram_xfer_time = self._histogram_cls(
3440
- name="vllm:nixl_xfer_time_seconds",
3441
- documentation="Histogram of transfer duration for NIXL KV Cache transfers.",
3442
- buckets=buckets[1:],
3443
- labelnames=labelnames,
3444
- )
3445
- self.nixl_histogram_xfer_time = create_metric_per_engine(
3446
- nixl_histogram_xfer_time, self.per_engine_labelvalues
3447
- )
3448
- nixl_histogram_post_time = self._histogram_cls(
3449
- name="vllm:nixl_post_time_seconds",
3450
- documentation="Histogram of transfer post time for NIXL KV"
3451
- " Cache transfers.",
3452
- buckets=buckets,
3453
- labelnames=labelnames,
3454
- )
3455
- self.nixl_histogram_post_time = create_metric_per_engine(
3456
- nixl_histogram_post_time, self.per_engine_labelvalues
3457
- )
3458
- # uniform 2kb to 16gb range
3459
- buckets = [2 ** (10 + i) for i in range(1, 25, 2)]
3460
- nixl_histogram_bytes_transferred = self._histogram_cls(
3461
- name="vllm:nixl_bytes_transferred",
3462
- documentation="Histogram of bytes transferred per NIXL KV Cache transfers.",
3463
- buckets=buckets,
3464
- labelnames=labelnames,
3465
- )
3466
- self.nixl_histogram_bytes_transferred = create_metric_per_engine(
3467
- nixl_histogram_bytes_transferred, self.per_engine_labelvalues
3468
- )
3469
- buckets = [
3470
- 10,
3471
- 20,
3472
- 30,
3473
- 50,
3474
- 75,
3475
- 100,
3476
- 200,
3477
- 400,
3478
- 1000,
3479
- 2000,
3480
- 4000,
3481
- 10000,
3482
- 20000,
3483
- 50000,
3484
- ]
3485
- nixl_histogram_num_descriptors = self._histogram_cls(
3486
- name="vllm:nixl_num_descriptors",
3487
- documentation="Histogram of number of descriptors per NIXL"
3488
- " KV Cache transfers.",
3489
- buckets=buckets,
3490
- labelnames=labelnames,
3491
- )
3492
- self.nixl_histogram_num_descriptors = create_metric_per_engine(
3493
- nixl_histogram_num_descriptors, self.per_engine_labelvalues
3494
- )
3495
- counter_nixl_num_failed_transfers = self._counter_cls(
3496
- name="vllm:nixl_num_failed_transfers",
3497
- documentation="Number of failed NIXL KV Cache transfers.",
3498
- labelnames=labelnames,
3499
- )
3500
- self.counter_nixl_num_failed_transfers = create_metric_per_engine(
3501
- counter_nixl_num_failed_transfers, self.per_engine_labelvalues
3502
- )
3503
- counter_nixl_num_failed_notifications = self._counter_cls(
3504
- name="vllm:nixl_num_failed_notifications",
3505
- documentation="Number of failed NIXL KV Cache notifications.",
3506
- labelnames=labelnames,
3507
- )
3508
- self.counter_nixl_num_failed_notifications = create_metric_per_engine(
3509
- counter_nixl_num_failed_notifications, self.per_engine_labelvalues
3510
- )
3511
-
3512
- counter_nixl_num_kv_expired_reqs = self._counter_cls(
3513
- name="vllm:nixl_num_kv_expired_reqs",
3514
- documentation="Number of requests that had their KV expire. "
3515
- "NOTE: This metric is tracked on the P instance.",
3516
- labelnames=labelnames,
3517
- )
3518
- self.counter_nixl_num_kv_expired_reqs = create_metric_per_engine(
3519
- counter_nixl_num_kv_expired_reqs, self.per_engine_labelvalues
3520
- )
3521
-
3522
- def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
3523
- for prom_obj, list_item_key in zip(
3524
- [
3525
- self.nixl_histogram_xfer_time,
3526
- self.nixl_histogram_post_time,
3527
- self.nixl_histogram_bytes_transferred,
3528
- self.nixl_histogram_num_descriptors,
3529
- ],
3530
- [
3531
- "transfer_duration",
3532
- "post_duration",
3533
- "bytes_transferred",
3534
- "num_descriptors",
3535
- ],
3536
- ):
3537
- for list_item in transfer_stats_data[list_item_key]:
3538
- prom_obj[engine_idx].observe(list_item)
3539
- for counter_obj, counter_item_key in zip(
3540
- [
3541
- self.counter_nixl_num_failed_transfers,
3542
- self.counter_nixl_num_failed_notifications,
3543
- self.counter_nixl_num_kv_expired_reqs,
3544
- ],
3545
- ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"],
3546
- ):
3547
- for list_item in transfer_stats_data[counter_item_key]:
3548
- counter_obj[engine_idx].inc(list_item)
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+1
-1
@@ -114,7 +114,7 @@ def derive_mamba_conv_split(
114
assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}"
115
116
# NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted
117
- # in nixl_connector __init__. Use it directly instead of heuristic detection.
117
+ # in nixl worker __init__. Use it directly instead of heuristic detection.
118
assert is_conv_state_dim_first(), "3-read requires DS conv state layout"
119
local_conv_dim = conv_shape[0] # DS: (conv_dim_local, conv_rows)
120
conv_rows = conv_shape[1]