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 7, 02:45 UTC
[Core][KVConnector] Support HMA+NixlConnector (#35758)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Mar 6, 2026 at 08:51 UTC
5b3ba94ab4bd9da739bcc27cdd05505467fa499e
10 files changed
+669
-230
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+9
@@ -12,6 +12,7 @@ tp_configs=(
12
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA case
13
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
14
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
15
+ "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" # SW model
16
)
17
dp_ep_configs=(
18
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
@@ -26,6 +27,14 @@ else
27
configs=("${tp_configs[@]}")
28
fi
29
30
+if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
31
+ # Append ENABLE_HMA_FLAG=1 to each config in the selected array
32
+ echo "ENABLE_HMA_FLAG is set, appending ENABLE_HMA_FLAG=1 to each config"
33
+ for i in "${!configs[@]}"; do
34
+ configs[$i]="ENABLE_HMA_FLAG=1 ${configs[$i]}"
35
+ done
36
+fi
37
+
38
run_tests() {
39
local label=$1
40
local extra_args=$2
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+36
-1
@@ -5,6 +5,12 @@ set -xe
5
KV_BUFFER_DEVICE="cuda" # Default to cuda
6
ATTENTION_BACKEND="" # Default to empty (use vllm default)
7
CROSS_LAYERS_BLOCKS="False"
8
+ENABLE_HMA_VAR="" # Default to empty (HMA disabled by default for kv connector)
9
+# Check for ENABLE_HMA_FLAG environment variable
10
+if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
11
+ ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
12
+fi
13
+
14
while [[ $# -gt 0 ]]; do
15
case $1 in
16
--kv_buffer_device)
@@ -31,6 +37,12 @@ echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"
37
if [[ -n "$ATTENTION_BACKEND" ]]; then
38
echo "Using attention backend: $ATTENTION_BACKEND"
39
fi
40
+if [[ -n "$ENABLE_HMA_VAR" ]]; then
41
+ echo "HMA (Hybrid KV Cache Manager) enabled"
42
+fi
43
+if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
44
+ echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
45
+fi
46
47
DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
48
if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
@@ -70,6 +82,8 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
82
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
83
PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
84
DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
85
+# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
86
+VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
87
88
# Find the git repository root directory
89
GIT_ROOT=$(git rev-parse --show-toplevel)
@@ -151,14 +165,24 @@ run_tests_for_model() {
165
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
166
--tensor-parallel-size $PREFILLER_TP_SIZE \
167
--kv-transfer-config '$KV_CONFIG'"
168
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
169
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
170
+ for arg in "${extra_args[@]}"; do
171
+ BASE_CMD="${BASE_CMD} $arg"
172
+ done
173
+ fi
174
175
# Add attention backend config if specified
176
if [[ -n "$ATTENTION_BACKEND" ]]; then
177
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
178
fi
179
180
+ # Add HMA flag if specified
181
+ if [[ -n "$ENABLE_HMA_VAR" ]]; then
182
+ BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
183
+ fi
184
+
185
FULL_CMD="$BASE_CMD"
161
-
186
eval "$FULL_CMD &"
187
188
# Store host and port for proxy configuration
@@ -193,12 +217,23 @@ run_tests_for_model() {
217
--block-size ${DECODE_BLOCK_SIZE} \
218
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
219
--kv-transfer-config '$KV_CONFIG'"
220
+ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
221
+ IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
222
+ for arg in "${extra_args[@]}"; do
223
+ BASE_CMD="${BASE_CMD} $arg"
224
+ done
225
+ fi
226
227
# Add attention backend config if specified
228
if [[ -n "$ATTENTION_BACKEND" ]]; then
229
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
230
fi
231
232
+ # Add HMA flag if specified
233
+ if [[ -n "$ENABLE_HMA_VAR" ]]; then
234
+ BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
235
+ fi
236
+
237
# DP-EP attention mode
238
if [[ -z "$DP_EP" ]]; then
239
BASE_CMD="${BASE_CMD} --tensor-parallel-size $DECODER_TP_SIZE"
tests/v1/kv_connector/nixl_integration/test_accuracy.py
+1
@@ -17,6 +17,7 @@ EXPECTED_VALUES = {
17
"deepseek-ai/deepseek-vl2-small": 0.59,
18
"deepseek-ai/deepseek-vl2-tiny": 0.19,
19
"deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,
20
+ "google/gemma-3-4b-it": 0.74,
21
}
22
23
SIMPLE_PROMPT = (
tests/v1/kv_connector/unit/test_nixl_connector.py
+122
-46
@@ -59,7 +59,12 @@ from vllm.v1.request import RequestStatus
59
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
60
from vllm.v1.worker.utils import AttentionGroup
61
62
-from .utils import create_request, create_scheduler, create_vllm_config
62
+from .utils import (
63
+ create_request,
64
+ create_scheduler,
65
+ create_vllm_config,
66
+ make_kv_cache_config,
67
+)
68
69
70
@pytest.fixture(scope="module", autouse=True)
@@ -263,7 +268,7 @@ def test_basic_interface():
268
req_meta = kv_connector_metadata.reqs_to_recv[request_id]
269
270
for block_id, block in zip(
266
- req_meta.local_block_ids,
271
+ req_meta.local_block_ids[0],
272
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
273
request_id
274
],
@@ -327,7 +332,9 @@ def test_kv_transfer_handshake(dist_init):
332
333
# Prefill connector will register KV cache to populate proper handshake
334
# metadata.
330
- prefill_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
335
+ prefill_connector = NixlConnector(
336
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
337
+ )
338
kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
339
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
340
)
@@ -367,13 +374,17 @@ def test_kv_transfer_handshake(dist_init):
374
do_remote_decode=True,
375
)
376
request.status = RequestStatus.FINISHED_LENGTH_CAPPED
370
- delay, kv_connector_metadata = scheduler.get_kv_connector().request_finished(
371
- request, [0, 1, 2]
377
+ delay, kv_connector_metadata = (
378
+ scheduler.get_kv_connector().request_finished_all_groups(
379
+ request, ([0, 1, 2],)
380
+ )
381
)
382
assert delay
383
384
# Decode connector will be able to create handshake with the prefill connector.
376
- decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
385
+ decode_connector = NixlConnector(
386
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
387
+ )
388
decode_connector.register_kv_caches(kv_caches)
389
390
# Here we are testing the retrieval of NIXLAgentMetadata.
@@ -404,9 +415,16 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
415
REMOTE_ENGINE_ID = "remote_engine"
416
417
def __init__(
407
- self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
418
+ self,
419
+ *args,
420
+ hand_shake_latency: float = 1.8,
421
+ kv_cache_layout="HND",
422
+ kv_cache_config=None,
423
+ **kwargs,
424
):
409
- super().__init__(*args, **kwargs)
425
+ if kv_cache_config is None:
426
+ kv_cache_config = make_kv_cache_config(block_size=16)
427
+ super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs)
428
self._hand_shake_latency = hand_shake_latency
429
self.kv_cache_layout = kv_cache_layout
430
# Mock register_kv_caches attribute needed for tests that do not call it.
@@ -507,7 +525,9 @@ class TestNixlHandshake:
525
request_id = "req_id"
526
527
# Test worker role in decode server.
510
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
528
+ connector = NixlConnector(
529
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
530
+ )
531
connector.connector_worker = FakeNixlConnectorWorker(
532
vllm_config, connector.engine_id, hand_shake_latency=0
533
)
@@ -528,13 +548,15 @@ class TestNixlHandshake:
548
num_xfers -= 1
549
metadata.add_new_req_to_recv(
550
request_id=request_id,
531
- local_block_ids=[num_xfers + 1, num_xfers + 2, num_xfers + 3],
551
+ local_block_ids=([num_xfers + 1, num_xfers + 2, num_xfers + 3],),
552
kv_transfer_params={
533
- "remote_block_ids": [
534
- num_xfers + 4,
535
- num_xfers + 5,
536
- num_xfers + 6,
537
- ],
553
+ "remote_block_ids": (
554
+ [
555
+ num_xfers + 4,
556
+ num_xfers + 5,
557
+ num_xfers + 6,
558
+ ],
559
+ ),
560
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
561
"remote_request_id": f"prefill-{request_id}",
562
"remote_host": "localhost",
@@ -594,16 +616,18 @@ class TestNixlHandshake:
616
vllm_config.parallel_config.tensor_parallel_size = decode_tp_size
617
618
# Test worker role in decode server.
597
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
619
+ connector = NixlConnector(
620
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
621
+ )
622
connector.connector_worker = FakeNixlConnectorWorker(
623
vllm_config, connector.engine_id
624
)
625
metadata = NixlConnectorMetadata()
626
metadata.add_new_req_to_recv(
627
request_id="id",
604
- local_block_ids=[1, 2, 3],
628
+ local_block_ids=([1, 2, 3],),
629
kv_transfer_params={
606
- "remote_block_ids": [4, 5, 6],
630
+ "remote_block_ids": ([4, 5, 6],),
631
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
632
"remote_request_id": "prefill-id",
633
"remote_host": "localhost",
@@ -652,7 +676,9 @@ class TestNixlHandshake:
676
local_tp_size = 1
677
vllm_config.parallel_config.tensor_parallel_size = local_tp_size
678
655
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
679
+ connector = NixlConnector(
680
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
681
+ )
682
connector.connector_worker = FakeNixlConnectorWorker(
683
vllm_config, connector.engine_id, hand_shake_latency=0
684
)
@@ -717,8 +743,12 @@ class TestNixlHandshake:
743
p_tp_size = 2
744
745
# Build two separate connectors/workers to emulate P TP=2 ranks.
720
- conn_p0 = NixlConnector(vllm_config, KVConnectorRole.WORKER)
721
- conn_p1 = NixlConnector(vllm_config, KVConnectorRole.WORKER)
746
+ conn_p0 = NixlConnector(
747
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
748
+ )
749
+ conn_p1 = NixlConnector(
750
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
751
+ )
752
conn_p0.connector_worker = FakeNixlConnectorWorker(
753
vllm_config, conn_p0.engine_id, hand_shake_latency=0
754
)
@@ -815,7 +845,9 @@ class TestNixlHandshake:
845
vllm_config = create_vllm_config()
846
847
# Test worker role in decode server.
818
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
848
+ connector = NixlConnector(
849
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
850
+ )
851
connector.connector_worker = FakeNixlConnectorWorker(
852
vllm_config, connector.engine_id
853
)
@@ -827,9 +859,9 @@ class TestNixlHandshake:
859
for i in range(total_reqs):
860
metadata.add_new_req_to_recv(
861
request_id=f"id_{i}",
830
- local_block_ids=[1, 2, 3],
862
+ local_block_ids=([1, 2, 3],),
863
kv_transfer_params={
832
- "remote_block_ids": [4, 5, 6],
864
+ "remote_block_ids": ([4, 5, 6],),
865
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
866
"remote_request_id": f"prefill-id-{i}",
867
"remote_host": "localhost",
@@ -884,7 +916,9 @@ class TestNixlHandshake:
916
return_value=2,
917
):
918
# Initialize connector and worker (with fake NIXL wrapper)
887
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
919
+ connector = NixlConnector(
920
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
921
+ )
922
connector.connector_worker = FakeNixlConnectorWorker(
923
vllm_config, connector.engine_id, hand_shake_latency=0
924
)
@@ -934,7 +968,9 @@ class TestNixlHandshake:
968
return_value=2,
969
):
970
# Initialize connector and worker (with fake NIXL wrapper)
937
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
971
+ connector = NixlConnector(
972
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
973
+ )
974
connector.connector_worker = FakeNixlConnectorWorker(
975
vllm_config,
976
connector.engine_id,
@@ -979,7 +1015,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1015
vllm_config = create_vllm_config()
1016
1017
# Test worker role in decode server.
982
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1018
+ connector = NixlConnector(
1019
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1020
+ )
1021
connector.connector_worker = FakeNixlConnectorWorker(
1022
vllm_config, connector.engine_id, hand_shake_latency=0
1023
)
@@ -993,9 +1031,9 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1031
metadata = NixlConnectorMetadata()
1032
metadata.add_new_req_to_recv(
1033
request_id=request_id,
996
- local_block_ids=[1, 2, 3],
1034
+ local_block_ids=([1, 2, 3],),
1035
kv_transfer_params={
998
- "remote_block_ids": [4, 5, 6],
1036
+ "remote_block_ids": ([4, 5, 6],),
1037
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1038
"remote_request_id": f"prefill-{request_id}",
1039
"remote_host": "localhost",
@@ -1448,7 +1486,9 @@ def test_register_kv_caches(
1486
mock_get_attn_backend.return_value = backend_cls
1487
1488
# Create connector
1451
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1489
+ connector = NixlConnector(
1490
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1491
+ )
1492
connector.connector_worker = FakeNixlConnectorWorker(
1493
vllm_config, connector.engine_id, hand_shake_latency=0
1494
)
@@ -1676,7 +1716,9 @@ def test_kv_buffer_to_nixl_memory_types(
1716
),
1717
): # noqa: E501
1718
# Create connector and replace its worker with a fake one for isolation
1679
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1719
+ connector = NixlConnector(
1720
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1721
+ )
1722
1723
# Verify get_reg_descs was called with the correct memory_type
1724
assert connector.connector_worker.kv_buffer_device == kv_buffer_device
@@ -1692,9 +1734,15 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
1734
vllm_config = create_vllm_config()
1735
1736
scheduler = NixlConnectorScheduler(
1695
- vllm_config, vllm_config.kv_transfer_config.engine_id
1737
+ vllm_config,
1738
+ vllm_config.kv_transfer_config.engine_id,
1739
+ make_kv_cache_config(block_size=16),
1740
+ )
1741
+ worker = NixlConnectorWorker(
1742
+ vllm_config,
1743
+ vllm_config.kv_transfer_config.engine_id,
1744
+ make_kv_cache_config(block_size=16),
1745
)
1697
- worker = NixlConnectorWorker(vllm_config, vllm_config.kv_transfer_config.engine_id)
1746
nixl_wrapper = worker.nixl_wrapper
1747
1748
with (
@@ -1756,7 +1804,9 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
1804
1805
scheduler = create_scheduler(vllm_config)
1806
# KVConnector Worker in P
1759
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1807
+ connector = NixlConnector(
1808
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1809
+ )
1810
connector.connector_worker = FakeNixlConnectorWorker(
1811
vllm_config, connector.engine_id, hand_shake_latency=0
1812
)
@@ -1875,12 +1925,14 @@ class FailingNixlWrapper(FakeNixlWrapper):
1925
("transfer_exception", {"fail_transfer_exception": True}, True),
1926
],
1927
)
1928
+@pytest.mark.parametrize("enable_hma", [False, True])
1929
def test_transfer_failure_logging(
1930
default_vllm_config,
1931
dist_init,
1932
failure_type,
1933
wrapper_config,
1934
needs_get_finished,
1935
+ enable_hma,
1936
):
1937
"""Test that transfer failures are logged with structured context.
1938
@@ -1897,9 +1949,16 @@ def test_transfer_failure_logging(
1949
1950
vllm_config = create_vllm_config()
1951
1900
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1952
+ connector = NixlConnector(
1953
+ vllm_config,
1954
+ KVConnectorRole.WORKER,
1955
+ make_kv_cache_config(block_size=16, hma_enabled=enable_hma),
1956
+ )
1957
connector.connector_worker = FakeNixlConnectorWorker(
1902
- vllm_config, connector.engine_id, hand_shake_latency=0.0
1958
+ vllm_config,
1959
+ connector.engine_id,
1960
+ hand_shake_latency=0.0,
1961
+ kv_cache_config=connector._kv_cache_config,
1962
)
1963
1964
# Configure FailingNixlWrapper to fail in the specified way
@@ -1910,8 +1969,17 @@ def test_transfer_failure_logging(
1969
1970
# For notification_failed, we need empty local blocks
1971
# (full cache hit path to trigger send_notif)
1913
- local_blocks = [] if failure_type == "notification_failed" else [10, 11, 12]
1914
- remote_blocks = [20, 21, 22]
1972
+ local_blocks: tuple[()] | tuple[list[int], ...]
1973
+ if enable_hma:
1974
+ # HMA enabled: multiple groups (FA + SW)
1975
+ local_blocks = (
1976
+ () if failure_type == "notification_failed" else ([10, 11, 12], [13, 14])
1977
+ )
1978
+ remote_blocks = [[20, 21, 22], [23, 24]]
1979
+ else:
1980
+ # HMA disabled: single group
1981
+ local_blocks = () if failure_type == "notification_failed" else ([10, 11, 12],)
1982
+ remote_blocks = [[20, 21, 22]]
1983
1984
metadata = NixlConnectorMetadata()
1985
metadata.add_new_req_to_recv(
@@ -2007,7 +2075,9 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2075
"""Test that handshake failures mark blocks invalid and return via get_finished."""
2076
vllm_config = create_vllm_config()
2077
2010
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
2078
+ connector = NixlConnector(
2079
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2080
+ )
2081
connector.connector_worker = FakeNixlConnectorWorker(
2082
vllm_config, connector.engine_id, hand_shake_latency=0.1
2083
)
@@ -2017,9 +2087,9 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2087
metadata = NixlConnectorMetadata()
2088
metadata.add_new_req_to_recv(
2089
request_id=request_id,
2020
- local_block_ids=[1, 2, 3],
2090
+ local_block_ids=([1, 2, 3],),
2091
kv_transfer_params={
2022
- "remote_block_ids": [4, 5, 6],
2092
+ "remote_block_ids": ([4, 5, 6],),
2093
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
2094
"remote_request_id": f"prefill-{request_id}",
2095
"remote_host": "localhost",
@@ -2058,7 +2128,9 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2128
and return via get_finished."""
2129
vllm_config = create_vllm_config()
2130
2061
- connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
2131
+ connector = NixlConnector(
2132
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2133
+ )
2134
connector.connector_worker = FakeNixlConnectorWorker(
2135
vllm_config, connector.engine_id, hand_shake_latency=0
2136
)
@@ -2068,9 +2140,9 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2140
metadata = NixlConnectorMetadata()
2141
metadata.add_new_req_to_recv(
2142
request_id=request_id,
2071
- local_block_ids=[7, 8, 9],
2143
+ local_block_ids=([7, 8, 9],),
2144
kv_transfer_params={
2073
- "remote_block_ids": [10, 11, 12],
2145
+ "remote_block_ids": ([10, 11, 12],),
2146
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
2147
"remote_request_id": f"prefill-{request_id}",
2148
"remote_host": "localhost",
@@ -2154,7 +2226,9 @@ def test_compatibility_hash_validation(
2226
"enforce_handshake_compat": enforce_handshake_compat
2227
},
2228
)
2157
- decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2229
+ decode_connector = NixlConnector(
2230
+ local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2231
+ )
2232
decode_worker = decode_connector.connector_worker
2233
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
2234
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
@@ -2267,7 +2341,9 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2341
model="facebook/opt-125m",
2342
block_size=16,
2343
)
2270
- decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
2344
+ decode_connector = NixlConnector(
2345
+ local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2346
+ )
2347
decode_worker = decode_connector.connector_worker
2348
2349
backend = get_current_attn_backend(local_vllm_config)
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
new
+203
@@ -0,0 +1,203 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for NixlConnectorScheduler sw_sizes calculation with HMA."""
4
+
5
+from unittest.mock import patch
6
+
7
+import pytest
8
+
9
+from vllm import LLM, SamplingParams
10
+from vllm.config import KVTransferConfig
11
+from vllm.v1.core.single_type_kv_cache_manager import (
12
+ FullAttentionManager,
13
+ SlidingWindowManager,
14
+)
15
+
16
+from .utils import (
17
+ create_vllm_config,
18
+ make_kv_cache_config,
19
+)
20
+
21
+
22
+@pytest.mark.cpu_test
23
+@pytest.mark.parametrize(
24
+ "hma_enabled,expected_sw_sizes",
25
+ [
26
+ # HMA enabled: FullAttentionSpec (0) + SlidingWindowSpec (2048/16=128)
27
+ (True, [0, 128 + 1]),
28
+ # HMA disabled: only FullAttentionSpec (0)
29
+ (False, [0]),
30
+ ],
31
+)
32
+@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")
33
+def test_sw_sizes(mock_platform, hma_enabled, expected_sw_sizes):
34
+ """Test sw_sizes is correctly computed based on HMA enabled/disabled."""
35
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
36
+ NixlConnectorScheduler,
37
+ )
38
+
39
+ mock_platform.device_type = "cpu"
40
+
41
+ block_size = 16
42
+ vllm_config = create_vllm_config(block_size=block_size)
43
+ # SW 2048 tokens=>128 blocks
44
+ kv_cache_config = make_kv_cache_config(
45
+ block_size=block_size, hma_enabled=hma_enabled, sw_size=2048
46
+ )
47
+
48
+ scheduler = NixlConnectorScheduler(
49
+ vllm_config=vllm_config,
50
+ engine_id="test-engine",
51
+ kv_cache_config=kv_cache_config,
52
+ )
53
+ # in number of blocks
54
+ assert scheduler.blocks_per_sw == expected_sw_sizes, (
55
+ f"Expected sw_sizes={expected_sw_sizes}, got {scheduler.blocks_per_sw}"
56
+ )
57
+
58
+
59
+@pytest.mark.cpu_test
60
+def test_logical_to_kernel_block_ids_with_hma():
61
+ """Test _logical_to_kernel_block_ids expands blocks when HMA is enabled.
62
+
63
+ When HMA is enabled, the logical block size may differ from the kernel
64
+ block size. Each logical block maps to multiple kernel blocks.
65
+ """
66
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
67
+ NixlConnectorWorker,
68
+ )
69
+
70
+ # Create a mock worker with just the required attributes
71
+ # (use __new__ to skip __init__)
72
+ worker = object.__new__(NixlConnectorWorker)
73
+
74
+ # Simulate HMA scenario: logical block size = 32, kernel block size = 16
75
+ # So each logical block maps to 2 kernel blocks eg [0]->[0,1]
76
+ worker._physical_blocks_per_logical_kv_block = 2
77
+
78
+ # Test conversion: FA + SW group
79
+ logical_block_ids = [[0, 1, 2], [3, 4]]
80
+ kernel_block_ids = worker._logical_to_kernel_block_ids(logical_block_ids)
81
+
82
+ expected_kernel_block_ids = [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]]
83
+ assert kernel_block_ids == expected_kernel_block_ids, (
84
+ f"Expected {expected_kernel_block_ids}, got {kernel_block_ids}"
85
+ )
86
+
87
+
88
+@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
89
+def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
90
+ """Test that a prefill instance returns fewer "remote blocks" for the SWA groups
91
+ when sequence exceeds the sliding window.
92
+ """
93
+ kv_transfer_config = KVTransferConfig(
94
+ kv_connector="NixlConnector",
95
+ kv_role="kv_both",
96
+ )
97
+ block_size = 16
98
+ llm_kwargs = {
99
+ "model": model_name,
100
+ "enforce_eager": True,
101
+ "gpu_memory_utilization": 0.5,
102
+ "kv_transfer_config": kv_transfer_config,
103
+ "max_model_len": 2048,
104
+ # NOTE: Make sure HMA is enabled
105
+ "disable_hybrid_kv_cache_manager": False,
106
+ "max_num_batched_tokens": 1024,
107
+ "enable_prefix_caching": False,
108
+ "block_size": block_size,
109
+ }
110
+
111
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
112
+
113
+ def run_hma_test(llm: LLM):
114
+ remote_prefill_opts = {
115
+ "do_remote_decode": True,
116
+ "do_remote_prefill": False,
117
+ "remote_engine_id": None,
118
+ "remote_block_ids": None,
119
+ "remote_host": None,
120
+ "remote_port": None,
121
+ }
122
+ # Simulate sidecar request
123
+ sampling_params = SamplingParams(
124
+ temperature=0.0,
125
+ max_tokens=1,
126
+ extra_args={"kv_transfer_params": remote_prefill_opts},
127
+ )
128
+ scheduler = llm.llm_engine.engine_core.engine_core.scheduler
129
+ kv_managers = scheduler.kv_cache_manager.coordinator.single_type_managers
130
+ # HMA enabled with FA + SWA groups
131
+ assert len(kv_managers) > 2
132
+ for kv_manager in kv_managers:
133
+ assert isinstance(kv_manager, (SlidingWindowManager, FullAttentionManager))
134
+ req_to_blocks = kv_managers[0].req_to_blocks
135
+ assert len(req_to_blocks) == 0
136
+
137
+ # Process some request with length exceeding the sliding window
138
+ outputs = llm.generate(["hi" * 1401], sampling_params)
139
+ kv_params = outputs[0].kv_transfer_params
140
+
141
+ # +1 to account for overlapping window across blocks.
142
+ expected_num_remote_blocks = sw_size // block_size + 1
143
+ remote_block_ids = kv_params["remote_block_ids"]
144
+ assert (
145
+ len(remote_block_ids[0])
146
+ == expected_num_remote_blocks
147
+ < len(remote_block_ids[-1])
148
+ )
149
+ for group_block_ids in remote_block_ids[:-1]:
150
+ assert len(group_block_ids) == expected_num_remote_blocks
151
+
152
+ def run_test_and_cleanup():
153
+ llm = LLM(**llm_kwargs)
154
+ try:
155
+ run_hma_test(llm)
156
+ finally:
157
+ llm.llm_engine.engine_core.shutdown()
158
+
159
+ run_test_and_cleanup()
160
+
161
+
162
+@pytest.mark.cpu_test
163
+def test_nixl_metadata_hma_block_ids_structure():
164
+ """
165
+ Test that NixlConnectorMetadata correctly stores block IDs for multiple
166
+ KV cache groups when HMA is enabled.
167
+ """
168
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
169
+ NixlConnectorMetadata,
170
+ )
171
+
172
+ metadata = NixlConnectorMetadata()
173
+
174
+ # Add request with block IDs for 2 groups (FA + SW)
175
+ fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7] # 8 blocks for FA
176
+ sw_blocks = [8, 9, 10, 11] # 4 blocks for SW (clipped)
177
+
178
+ metadata.add_new_req_to_recv(
179
+ request_id="test-req-hma",
180
+ local_block_ids=(fa_blocks, sw_blocks),
181
+ kv_transfer_params={
182
+ "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [18, 19, 20, 21]),
183
+ "remote_engine_id": "remote-engine",
184
+ "remote_request_id": "prefill-test-req-hma",
185
+ "remote_host": "localhost",
186
+ "remote_port": 1234,
187
+ "tp_size": 1,
188
+ },
189
+ )
190
+
191
+ assert "test-req-hma" in metadata.reqs_to_recv
192
+ req_meta = metadata.reqs_to_recv["test-req-hma"]
193
+
194
+ # Verify local block IDs structure
195
+ assert len(req_meta.local_block_ids) == 2
196
+ assert list(req_meta.local_block_ids[0]) == fa_blocks
197
+ assert list(req_meta.local_block_ids[1]) == sw_blocks
198
+
199
+ # Verify remote block IDs structure
200
+ assert req_meta.remote is not None
201
+ assert len(req_meta.remote.block_ids) == 2
202
+ assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
203
+ assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+3
-1
@@ -208,7 +208,9 @@ def test_prefix_cache_lifecycle():
208
209
# Ensure we send all block ids, including the partial blocks,
210
# even if there is a cache hit.
211
- assert len(kv_transfer_params["remote_block_ids"]) == (NUM_EXTERNAL_FULL_BLOCKS + 1)
211
+ # remote_block_ids is BlockIds (tuple of lists); sum block counts across groups.
212
+ num_remote_blocks = sum(len(g) for g in kv_transfer_params["remote_block_ids"])
213
+ assert num_remote_blocks == (NUM_EXTERNAL_FULL_BLOCKS + 1)
214
215
# STEP (2): Ensure it is freed.
216
scheduler_output = scheduler.schedule()
tests/v1/kv_connector/unit/utils.py
+53
-15
@@ -36,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
36
FullAttentionSpec,
37
KVCacheConfig,
38
KVCacheGroupSpec,
39
+ SlidingWindowSpec,
40
)
41
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
42
from vllm.v1.request import Request
@@ -142,24 +143,26 @@ def create_vllm_config(
143
def create_scheduler(
144
vllm_config: VllmConfig,
145
num_blocks: int = 10000,
146
+ kv_cache_config: KVCacheConfig | None = None,
147
) -> Scheduler:
148
"""Initialize Scheduler For Testing."""
149
block_size = vllm_config.cache_config.block_size
148
- kv_cache_config = KVCacheConfig(
149
- num_blocks=num_blocks, # A large number of blocks to hold all requests
150
- kv_cache_tensors=[],
151
- kv_cache_groups=[
152
- KVCacheGroupSpec(
153
- ["layer"],
154
- FullAttentionSpec(
155
- block_size=block_size,
156
- num_kv_heads=1,
157
- head_size=1,
158
- dtype=torch.float32,
159
- ),
160
- )
161
- ],
162
- )
150
+ if kv_cache_config is None:
151
+ kv_cache_config = KVCacheConfig(
152
+ num_blocks=num_blocks, # A large number of blocks to hold all requests
153
+ kv_cache_tensors=[],
154
+ kv_cache_groups=[
155
+ KVCacheGroupSpec(
156
+ ["layer"],
157
+ FullAttentionSpec(
158
+ block_size=block_size,
159
+ num_kv_heads=1,
160
+ head_size=1,
161
+ dtype=torch.float32,
162
+ ),
163
+ )
164
+ ],
165
+ )
166
vllm_config.cache_config.num_gpu_blocks = num_blocks
167
return Scheduler(
168
vllm_config=vllm_config,
@@ -412,3 +415,38 @@ KVConnectorFactory.register_connector(
415
KVConnectorFactory.register_connector(
416
"MockKVConnector", __name__, MockKVConnector.__name__
417
)
418
+
419
+
420
+def make_kv_cache_config(
421
+ block_size: int,
422
+ hma_enabled: bool = False,
423
+ sw_size: int = 128,
424
+ num_blocks: int = 100,
425
+) -> KVCacheConfig:
426
+ kv_cache_groups = [
427
+ KVCacheGroupSpec(
428
+ ["layer0", "layer2"],
429
+ FullAttentionSpec(
430
+ block_size=block_size,
431
+ num_kv_heads=4,
432
+ head_size=16,
433
+ dtype=torch.float16,
434
+ ),
435
+ )
436
+ ]
437
+ if hma_enabled:
438
+ kv_cache_groups.append(
439
+ KVCacheGroupSpec(
440
+ ["layer1", "layer3"],
441
+ SlidingWindowSpec(
442
+ block_size=block_size,
443
+ num_kv_heads=4,
444
+ head_size=16,
445
+ dtype=torch.float16,
446
+ sliding_window=sw_size,
447
+ ),
448
+ )
449
+ )
450
+ return KVCacheConfig(
451
+ num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
452
+ )
vllm/distributed/kv_transfer/kv_connector/utils.py
+3
@@ -24,6 +24,9 @@ if TYPE_CHECKING:
24
logger = init_logger(__name__)
25
26
EngineId = str
27
+# block ids as returned by the hybrid KV cache manager. list[list[int]] are allow
28
+# mutability and are for connector internal use only.
29
+BlockIds = tuple[list[int], ...] | list[list[int]]
30
31
32
def get_kv_connector_cache_layout():
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+227
-167
@@ -3,7 +3,6 @@
3
import contextlib
4
import copy
5
import logging
6
-import math
6
import os
7
import queue
8
import sys
@@ -24,6 +23,7 @@ import zmq
23
from vllm import envs
24
from vllm.config import VllmConfig
25
from vllm.distributed.kv_transfer.kv_connector.utils import (
26
+ BlockIds,
27
EngineId,
28
TpKVTopology,
29
get_current_attn_backend,
@@ -38,6 +38,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
38
KVConnectorHandshakeMetadata,
39
KVConnectorMetadata,
40
KVConnectorRole,
41
+ SupportsHMA,
42
)
43
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
44
KVConnectorPromMetrics,
@@ -53,10 +54,12 @@ from vllm.distributed.parallel_state import (
54
from vllm.forward_context import ForwardContext
55
from vllm.logger import init_logger
56
from vllm.platforms import current_platform
57
+from vllm.utils.math_utils import cdiv
58
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
59
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
60
from vllm.v1.attention.backends.utils import get_kv_cache_layout
61
from vllm.v1.core.sched.output import SchedulerOutput
62
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec
63
from vllm.v1.worker.block_table import BlockTable
64
65
if TYPE_CHECKING:
@@ -205,6 +208,7 @@ def compute_nixl_compatibility_hash(
208
209
model_config = vllm_config.model_config
210
cache_config = vllm_config.cache_config
211
+ is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
212
213
factors = {
214
# Version compatibility
@@ -220,6 +224,7 @@ def compute_nixl_compatibility_hash(
224
"attn_backend_name": attn_backend_name,
225
"cache_dtype": str(cache_config.cache_dtype),
226
"cross_layers_blocks": cross_layers_blocks,
227
+ "is_hma_enabled": is_hma_enabled,
228
}
229
230
compat_hash = hash_factors(factors)
@@ -238,7 +243,7 @@ def compute_nixl_compatibility_hash(
243
244
@dataclass
245
class RemoteMeta:
241
- block_ids: list[int]
246
+ block_ids: BlockIds
247
host: str
248
port: int
249
engine_id: str
@@ -247,9 +252,9 @@ class RemoteMeta:
252
253
@dataclass
254
class ReqMeta:
250
- local_block_ids: list[int]
255
+ local_block_ids: BlockIds
256
# To be used when logical block size does not match the kernel block size
252
- local_physical_block_ids: list[int]
257
+ local_physical_block_ids: BlockIds
258
tp_size: int
259
remote: RemoteMeta | None = None
260
@@ -264,7 +269,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
269
270
def _add_new_req(
271
self,
267
- local_block_ids: list[int],
272
+ local_block_ids: BlockIds,
273
kv_transfer_params: dict[str, Any],
274
) -> ReqMeta:
275
return ReqMeta(
@@ -277,7 +282,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
282
def add_new_req_to_save(
283
self,
284
request_id: ReqId,
280
- local_block_ids: list[int],
285
+ local_block_ids: BlockIds,
286
kv_transfer_params: dict[str, Any],
287
):
288
self.reqs_to_save[request_id] = self._add_new_req(
@@ -287,7 +292,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
292
def add_new_req_to_recv(
293
self,
294
request_id: ReqId,
290
- local_block_ids: list[int],
295
+ local_block_ids: BlockIds,
296
kv_transfer_params: dict[str, Any],
297
):
298
req = self._add_new_req(local_block_ids, kv_transfer_params)
@@ -301,7 +306,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
306
self.reqs_to_recv[request_id] = req
307
308
304
-class NixlConnector(KVConnectorBase_V1):
309
+class NixlConnector(KVConnectorBase_V1, SupportsHMA):
310
@property
311
def prefer_cross_layer_blocks(self) -> bool:
312
backend = get_current_attn_backend(self._vllm_config)
@@ -326,22 +331,27 @@ class NixlConnector(KVConnectorBase_V1):
331
self,
332
vllm_config: VllmConfig,
333
role: KVConnectorRole,
329
- kv_cache_config: "KVCacheConfig | None" = None,
334
+ kv_cache_config: "KVCacheConfig",
335
):
336
super().__init__(vllm_config, role, kv_cache_config)
337
338
assert vllm_config.kv_transfer_config is not None
339
assert vllm_config.kv_transfer_config.engine_id is not None
340
+ for group in kv_cache_config.kv_cache_groups:
341
+ if isinstance(group.kv_cache_spec, MambaSpec):
342
+ raise ValueError("NixlConnector does not support Mamba models.")
343
self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
344
self.kv_transfer_config = vllm_config.kv_transfer_config
345
if role == KVConnectorRole.SCHEDULER:
346
self.connector_scheduler: NixlConnectorScheduler | None = (
339
- NixlConnectorScheduler(vllm_config, self.engine_id)
347
+ NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
348
)
349
self.connector_worker: NixlConnectorWorker | None = None
350
elif role == KVConnectorRole.WORKER:
351
self.connector_scheduler = None
344
- self.connector_worker = NixlConnectorWorker(vllm_config, self.engine_id)
352
+ self.connector_worker = NixlConnectorWorker(
353
+ vllm_config, self.engine_id, kv_cache_config
354
+ )
355
356
############################################################
357
# Class Methods
@@ -392,10 +402,10 @@ class NixlConnector(KVConnectorBase_V1):
402
assert self.connector_scheduler is not None
403
return self.connector_scheduler.build_connector_meta(scheduler_output)
404
395
- def request_finished(
405
+ def request_finished_all_groups(
406
self,
407
request: "Request",
398
- block_ids: list[int],
408
+ block_ids: tuple[list[int], ...],
409
) -> tuple[bool, dict[str, Any] | None]:
410
assert self.connector_scheduler is not None
411
return self.connector_scheduler.request_finished(request, block_ids)
@@ -518,10 +528,13 @@ class NixlConnector(KVConnectorBase_V1):
528
class NixlConnectorScheduler:
529
"""Implementation of Scheduler side methods"""
530
521
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
531
+ def __init__(
532
+ self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig"
533
+ ):
534
self.vllm_config = vllm_config
535
self.block_size = vllm_config.cache_config.block_size
536
self.engine_id: EngineId = engine_id
537
+ self.kv_cache_config = kv_cache_config
538
self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
539
self.side_channel_port = (
540
envs.VLLM_NIXL_SIDE_CHANNEL_PORT
@@ -534,8 +547,18 @@ class NixlConnectorScheduler:
547
self.use_host_buffer = (
548
vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
549
)
550
+ self._is_hma_required = (
551
+ not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
552
+ # Also handle unlikely SW-only model case instead of checking num_groups>1.
553
+ and any(
554
+ not isinstance(g.kv_cache_spec, FullAttentionSpec)
555
+ for g in kv_cache_config.kv_cache_groups
556
+ )
557
+ )
558
559
logger.info("Initializing NIXL Scheduler %s", engine_id)
560
+ if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
561
+ logger.info("Hybrid Memory Allocator is enabled with NIXL")
562
563
# Background thread for handling new handshake requests.
564
self._nixl_handshake_listener_t: threading.Thread | None = None
@@ -545,7 +568,7 @@ class NixlConnectorScheduler:
568
# Requests that need to start recv/send.
569
# New requests are added by update_state_after_alloc in
570
# the scheduler. Used to make metadata passed to Worker.
548
- self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
571
+ self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
572
self._reqs_need_save: dict[ReqId, Request] = {}
573
# Reqs to send and their expiration time
574
self._reqs_need_send: dict[ReqId, float] = {}
@@ -554,12 +577,54 @@ class NixlConnectorScheduler:
577
# remote prefill or aborted.
578
self._reqs_not_processed: set[ReqId] = set()
579
580
+ # Gather Sliding Window sizes for each kv cache group (if any) in number of
581
+ # blocks per KV cache group. This is used to clip the local attention window.
582
+ sw_sizes_tokens: list[tuple[int, int]] = [
583
+ (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
584
+ if isinstance(g.kv_cache_spec, SlidingWindowSpec)
585
+ else (0, self.block_size)
586
+ for g in kv_cache_config.kv_cache_groups
587
+ ]
588
+ # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
589
+ # account for boundary overlap eg window isn't fully aligned with blocks.
590
+ self.blocks_per_sw = [
591
+ cdiv(n_tokens, block_size) + 1 if n_tokens else 0
592
+ for n_tokens, block_size in sw_sizes_tokens
593
+ ]
594
+
595
def shutdown(self):
596
self._stop_event.set()
597
if self._nixl_handshake_listener_t is not None:
598
self._nixl_handshake_listener_t.join()
599
self._nixl_handshake_listener_t = None
600
601
+ def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
602
+ """
603
+ Clip the number of blocks to the sliding window size for each kv cache group
604
+ that employs SWA.
605
+ This is necessary because the KV Cache manager initially allocates blocks for
606
+ the entire sequence length, and successively cleans up blocks that are outside
607
+ the window prior to the `request_finished_all_groups` hook.
608
+ """
609
+ if len(block_ids) == 0 or not self._is_hma_required:
610
+ # No blocks to clip eg Full prefix cache hit or not a hybrid model.
611
+ return block_ids
612
+ # NOTE (NickLucche) This logic is currently handled at the connector level
613
+ # because offloading connectors might want to receive the whole sequence even
614
+ # for SWA groups. We will abstract this logic once the interface is more stable
615
+ assert len(block_ids) == len(self.blocks_per_sw), (
616
+ "Number of KV cache groups must match"
617
+ )
618
+ # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
619
+ return tuple(
620
+ [
621
+ blocks[-self.blocks_per_sw[i] :]
622
+ if self.blocks_per_sw[i] > 0
623
+ else blocks
624
+ for i, blocks in enumerate(block_ids)
625
+ ]
626
+ )
627
+
628
def set_xfer_handshake_metadata(
629
self, metadata: dict[int, KVConnectorHandshakeMetadata]
630
) -> None:
@@ -707,12 +772,18 @@ class NixlConnectorScheduler:
772
# If remote_blocks and num_external_tokens = 0, we have
773
# a full prefix cache hit on the D worker. We need to call
774
# send_notif in _read_blocks to free the memory on the P.
710
- local_block_ids = (
711
- blocks.get_unhashed_block_ids()
775
+
776
+ unhashed_local_block_ids: BlockIds = (
777
+ blocks.get_unhashed_block_ids_all_groups()
778
if num_external_tokens > 0
713
- else []
779
+ else ()
780
)
715
- # Get unhashed blocks to pull from remote.
781
+ local_block_ids = self.get_sw_clipped_blocks(
782
+ unhashed_local_block_ids
783
+ )
784
+
785
+ # Get unhashed blocks to pull from remote. Mind that a full prefix
786
+ # cache hit is indicated with an empty list.
787
self._reqs_need_recv[request.request_id] = (
788
request,
789
local_block_ids,
@@ -753,9 +824,10 @@ class NixlConnectorScheduler:
824
req = req_to_save
825
826
assert req.kv_transfer_params is not None
827
+ clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
828
meta.add_new_req_to_save(
829
request_id=req_id,
758
- local_block_ids=new_block_id_groups[0],
830
+ local_block_ids=clipped_block_id_groups,
831
kv_transfer_params=req.kv_transfer_params,
832
)
833
assert scheduler_output.num_scheduled_tokens is not None
@@ -786,7 +858,7 @@ class NixlConnectorScheduler:
858
def request_finished(
859
self,
860
request: "Request",
789
- block_ids: list[int],
861
+ block_ids: BlockIds,
862
) -> tuple[bool, dict[str, Any] | None]:
863
"""
864
Once a request is finished, determine whether request blocks
@@ -828,7 +900,7 @@ class NixlConnectorScheduler:
900
901
# TODO: check whether block_ids actually ever be 0. If not we could
902
# remove the conditional below
831
- delay_free_blocks = len(block_ids) > 0
903
+ delay_free_blocks = any(len(group) > 0 for group in block_ids)
904
905
if delay_free_blocks:
906
# Prefill request on remote. It will be read from D upon completion
@@ -841,6 +913,11 @@ class NixlConnectorScheduler:
913
self._reqs_need_send[request.request_id] = (
914
time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
915
)
916
+ # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
917
+ # trimming down after allocating for the whole sequence length. Empty
918
+ # blocks are always at the start of the list.
919
+ # Here we "unpad" blocks to send the actual remote blocks to be read.
920
+ block_ids = self.get_sw_clipped_blocks(block_ids)
921
922
return delay_free_blocks, dict(
923
do_remote_prefill=True,
@@ -857,7 +934,9 @@ class NixlConnectorScheduler:
934
class NixlConnectorWorker:
935
"""Implementation of Worker side methods"""
936
860
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
937
+ def __init__(
938
+ self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig"
939
+ ):
940
if NixlWrapper is None:
941
logger.error("NIXL is not available")
942
raise RuntimeError("NIXL is not available")
@@ -875,6 +954,14 @@ class NixlConnectorWorker:
954
self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(
955
"backends", ["UCX"]
956
)
957
+ self._is_hma_required = (
958
+ not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
959
+ and any(
960
+ not isinstance(g.kv_cache_spec, FullAttentionSpec)
961
+ for g in kv_cache_config.kv_cache_groups
962
+ )
963
+ )
964
+ self.kv_cache_config = kv_cache_config
965
966
# Agent.
967
non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
@@ -1017,10 +1104,6 @@ class NixlConnectorWorker:
1104
self.model_config = vllm_config.model_config
1105
self.cache_config = vllm_config.cache_config
1106
1020
- # TODO(mgoin): remove this once we have hybrid memory allocator
1021
- # Optimization for models with local attention (Llama 4)
1022
- # List of block window sizes for each layer for local attention
1023
- self.block_window_per_layer: list[int | None] = []
1107
self.use_mla = self.model_config.use_mla
1108
1109
# Get the attention backend from the first layer
@@ -1030,8 +1113,8 @@ class NixlConnectorWorker:
1113
self.backend_name = self.attn_backend.get_name()
1114
self.kv_cache_layout = get_kv_cache_layout()
1115
self.host_buffer_kv_cache_layout = self.kv_cache_layout
1033
- logger.debug("Detected attention backend %s", self.backend_name)
1034
- logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
1116
+ logger.info("Detected attention backend %s", self.backend_name)
1117
+ logger.info("Detected kv cache layout %s", self.kv_cache_layout)
1118
1119
# lazy initialized in register_kv_caches
1120
self.compat_hash: str | None = None
@@ -1238,9 +1321,15 @@ class NixlConnectorWorker:
1321
"remote_request_id": meta.remote.request_id,
1322
"remote_host": meta.remote.host,
1323
"remote_port": meta.remote.port,
1241
- "num_local_blocks": len(meta.local_block_ids),
1242
- "num_remote_blocks": len(meta.remote.block_ids),
1243
- "local_block_ids_sample": meta.local_block_ids[:10],
1324
+ "num_local_blocks": sum(
1325
+ len(group) for group in meta.local_block_ids
1326
+ ),
1327
+ "num_remote_blocks": sum(
1328
+ len(group) for group in meta.remote.block_ids
1329
+ ),
1330
+ "local_block_ids_sample": meta.local_block_ids[0][:10]
1331
+ if meta.local_block_ids
1332
+ else [],
1333
}
1334
)
1335
@@ -1301,8 +1390,10 @@ class NixlConnectorWorker:
1390
error=e,
1391
meta=meta,
1392
)
1304
- if req_meta := self._recving_metadata.get(req_id):
1305
- self._invalid_block_ids.update(req_meta.local_block_ids)
1393
+ if (
1394
+ req_meta := self._recving_metadata.get(req_id)
1395
+ ) and not self._is_hma_required:
1396
+ self._invalid_block_ids.update(req_meta.local_block_ids[0])
1397
self._failed_recv_reqs.add(req_id)
1398
1399
fut.add_done_callback(request_ready)
@@ -1370,6 +1461,10 @@ class NixlConnectorWorker:
1461
for cache in cache_list:
1462
base_addr = cache.data_ptr()
1463
if base_addr in seen_base_addresses:
1464
+ # NOTE (NickLucche) HMA employs memory pooling to share tensors
1465
+ # across groups. This results in skipping all tensors but the ones
1466
+ # pointed to by group0. Also, generally we will have more blocks
1467
+ # per tensor but fewer regions.
1468
continue
1469
1470
logger.debug(
@@ -1457,28 +1552,6 @@ class NixlConnectorWorker:
1552
self.register_local_xfer_handler(self.block_size)
1553
)
1554
1460
- # TODO(mgoin): Hybrid memory allocator is currently disabled for
1461
- # models with local attention (Llama 4). Can remove this once enabled.
1462
- if self.model_config.hf_config.model_type == "llama4":
1463
- from transformers import Llama4TextConfig
1464
-
1465
- assert isinstance(self.model_config.hf_text_config, Llama4TextConfig)
1466
- llama4_config = self.model_config.hf_text_config
1467
- no_rope_layers = llama4_config.no_rope_layers
1468
- chunk_size = llama4_config.attention_chunk_size
1469
- chunk_block_size = math.ceil(chunk_size / self.block_size)
1470
- for layer_idx in range(self.num_layers):
1471
- # no_rope_layers[layer_idx] == 0 means NoPE (global)
1472
- # Any other value means RoPE (local chunked)
1473
- is_local_attention = no_rope_layers[layer_idx] != 0
1474
- block_window = chunk_block_size if is_local_attention else None
1475
- self.block_window_per_layer.append(block_window)
1476
- logger.debug(
1477
- "Llama 4 block window per layer mapping: %s",
1478
- self.block_window_per_layer,
1479
- )
1480
- assert len(self.block_window_per_layer) == self.num_layers
1481
-
1555
# After KV Caches registered, listen for new connections.
1556
agent_metadata = NixlAgentMetadata(
1557
engine_id=self.engine_id,
@@ -1767,6 +1840,11 @@ class NixlConnectorWorker:
1840
# Num kv_heads > tp_size and P TP > D TP case, not supported
1841
assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id))
1842
1843
+ if self._is_hma_required:
1844
+ assert block_size_ratio == 1, (
1845
+ "HMA does not support different remote block size yet"
1846
+ )
1847
+
1848
kv_cache_layout = (
1849
self.kv_cache_layout
1850
if not self.use_host_buffer
@@ -1781,6 +1859,9 @@ class NixlConnectorWorker:
1859
"Remote is HND and local is NHD, enabled additional permute "
1860
"on local device KV."
1861
)
1862
+ assert not self._is_hma_required, (
1863
+ "HMA does not support block size post processing"
1864
+ )
1865
self.enable_permute_local_kv = True
1866
else:
1867
raise RuntimeError(
@@ -1836,13 +1917,15 @@ class NixlConnectorWorker:
1917
assert self.copy_blocks is not None
1918
1919
local_block_ids = meta.local_physical_block_ids
1839
- self.copy_blocks(
1840
- self.host_xfer_buffers,
1841
- self.device_kv_caches,
1842
- local_block_ids,
1843
- local_block_ids,
1844
- "h2d",
1845
- )
1920
+ # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups
1921
+ for group_block_ids in local_block_ids:
1922
+ self.copy_blocks(
1923
+ self.host_xfer_buffers,
1924
+ self.device_kv_caches,
1925
+ group_block_ids,
1926
+ group_block_ids,
1927
+ "h2d",
1928
+ )
1929
if logger.isEnabledFor(logging.DEBUG):
1930
logger.debug(
1931
"synced recved kv of request[%s] to device kv buffer,"
@@ -1868,13 +1951,14 @@ class NixlConnectorWorker:
1951
",".join(map(str, meta.local_physical_block_ids)),
1952
)
1953
# blocking
1871
- self.copy_blocks(
1872
- self.device_kv_caches,
1873
- self.host_xfer_buffers,
1874
- meta.local_physical_block_ids,
1875
- meta.local_physical_block_ids,
1876
- "d2h",
1877
- )
1954
+ for group_block_ids in meta.local_physical_block_ids:
1955
+ self.copy_blocks(
1956
+ self.device_kv_caches,
1957
+ self.host_xfer_buffers,
1958
+ group_block_ids,
1959
+ group_block_ids,
1960
+ "d2h",
1961
+ )
1962
1963
def post_process_device_kv_on_receive(
1964
self,
@@ -1973,8 +2057,9 @@ class NixlConnectorWorker:
2057
if not self.use_mla and (
2058
block_size_ratio > 1 or self.enable_permute_local_kv
2059
):
2060
+ assert not self._is_hma_required
2061
block_ids_for_blocksize_post_process[block_size_ratio].append(
1977
- meta.local_physical_block_ids
2062
+ meta.local_physical_block_ids[0]
2063
)
2064
for (
2065
block_size_ratio,
@@ -2106,8 +2191,9 @@ class NixlConnectorWorker:
2191
handle: The transfer handle.
2192
"""
2193
# Use .get() here as the metadata cleanup is handled by get_finished()
2109
- if meta := self._recving_metadata.get(req_id):
2110
- self._invalid_block_ids.update(meta.local_block_ids)
2194
+ # TODO (NickLucche) handle failed transfer for HMA.
2195
+ if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:
2196
+ self._invalid_block_ids.update(meta.local_block_ids[0])
2197
self.nixl_wrapper.release_xfer_handle(handle)
2198
self.xfer_stats.record_failed_transfer()
2199
@@ -2230,8 +2316,8 @@ class NixlConnectorWorker:
2316
2317
def _read_blocks(
2318
self,
2233
- local_block_ids: list[int],
2234
- remote_block_ids: list[int],
2319
+ local_block_ids: BlockIds,
2320
+ remote_block_ids: BlockIds,
2321
dst_engine_id: str,
2322
request_id: str,
2323
remote_request_id: str,
@@ -2246,22 +2332,30 @@ class NixlConnectorWorker:
2332
assert self.kv_topo is not None
2333
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
2334
if block_size_ratio > 1:
2249
- local_block_ids = self.get_mapped_blocks(
2250
- np.asarray(local_block_ids), block_size_ratio
2251
- )
2252
- if len(local_block_ids) > len(remote_block_ids):
2335
+ # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
2336
+ assert not self._is_hma_required
2337
+ local_block_ids0 = local_block_ids[0] if local_block_ids else []
2338
+ remote_block_ids0 = remote_block_ids[0]
2339
+ local_block_ids_mapped = self.get_mapped_blocks(
2340
+ np.asarray(local_block_ids0), block_size_ratio
2341
+ ).tolist()
2342
+ if len(local_block_ids_mapped) > len(remote_block_ids0):
2343
# NOTE:
2344
# get_mapped_blocks will always expand block_ids for n times.
2345
# ex:
2346
# prefill block_ids with block_size as 4:
2347
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2348
# Local decode block_ids with block_size as 16: [1, 2, 3]
2259
- # expland ecode block_ids with get_mapped_blocks from [1, 2, 3] to
2349
+ # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
2350
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
2351
# Then we clip local to align with prefill
2352
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
2353
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2264
- local_block_ids = local_block_ids[: len(remote_block_ids)]
2354
+ local_block_ids_mapped = local_block_ids_mapped[
2355
+ : len(remote_block_ids0)
2356
+ ]
2357
+ local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
2358
+ remote_block_ids = [remote_block_ids0]
2359
# NOTE(rob): having the staging blocks be on the READER side is
2360
# not going to work well (since we will have to call rearrange tensors).
2361
# after we detect the txn is complete (which means we cannot make the
@@ -2269,8 +2363,7 @@ class NixlConnectorWorker:
2363
# then we will need to have the staging blocks on the remote side.
2364
2365
# NOTE(rob): according to nvidia the staging blocks are used to
2272
- # saturate IB with heterogeneous TP sizes. We should remove the staging
2273
- # blocks until we are ready.
2366
+ # saturate IB with heterogeneous TP sizes.
2367
2368
# Number of D TP workers that will read from dst P. Propagate info
2369
# on notification so that dst worker can wait before freeing blocks.
@@ -2278,8 +2371,8 @@ class NixlConnectorWorker:
2371
2372
# Full prefix cache hit: do not need to read remote blocks,
2373
# just notify P worker that we have the blocks we need.
2281
- num_local_blocks = len(local_block_ids)
2282
- if num_local_blocks == 0:
2374
+ if len(local_block_ids) == 0:
2375
+ # A full prefix cache hit is indicated with an empty list.
2376
agent_name = self._remote_agents[dst_engine_id][remote_rank]
2377
try:
2378
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
@@ -2297,66 +2390,34 @@ class NixlConnectorWorker:
2390
self.xfer_stats.record_failed_notification()
2391
return
2392
2300
- # Partial prefix cache hit: just read uncomputed blocks.
2301
- num_remote_blocks = len(remote_block_ids)
2302
- assert num_local_blocks <= num_remote_blocks
2303
- if num_local_blocks < num_remote_blocks:
2304
- remote_block_ids = remote_block_ids[-num_local_blocks:]
2393
+ assert (
2394
+ len(remote_block_ids)
2395
+ == len(local_block_ids)
2396
+ == len(self.kv_cache_config.kv_cache_groups)
2397
+ )
2398
+ remote_block_ids = list(remote_block_ids)
2399
+ for i, remote_group in enumerate(remote_block_ids):
2400
+ num_remote_blocks = len(remote_group)
2401
+ num_local_blocks = len(local_block_ids[i])
2402
+ assert num_local_blocks <= num_remote_blocks
2403
+ # Partial prefix cache hit: just read uncomputed blocks.
2404
+ if num_local_blocks < num_remote_blocks:
2405
+ remote_block_ids[i] = remote_group[-num_local_blocks:]
2406
2407
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
2408
# corresponding rank. With heterogeneous TP, fixing D>P, the D tp
2409
# workers will issue xfers to parts of the P worker remote kv caches.
2410
2411
# Get descs ids.
2311
- local_block_descs_ids: np.ndarray
2312
- remote_block_descs_ids: np.ndarray
2313
-
2314
- if not self.block_window_per_layer:
2315
- # Default case: assume global attention
2316
- remote_block_descs_ids = self._get_block_descs_ids(
2317
- dst_engine_id,
2318
- remote_block_ids,
2319
- )
2320
- local_block_descs_ids = self._get_block_descs_ids(
2321
- self.engine_id,
2322
- local_block_ids,
2323
- block_size_ratio=block_size_ratio,
2324
- )
2325
- else:
2326
- # TODO(mgoin): remove this once we have hybrid memory allocator
2327
- # Optimization for models with local attention (Llama 4)
2328
- local_descs_list = []
2329
- remote_descs_list = []
2330
- for layer_idx, block_window in enumerate(self.block_window_per_layer):
2331
- # For each layer:
2332
- if block_window is None:
2333
- # If not chunked, we just use the
2334
- # full block lists (global attention)
2335
- layer_local_block_ids = local_block_ids
2336
- layer_remote_block_ids = remote_block_ids
2337
- else:
2338
- # If chunked, get the last block_window blocks
2339
- layer_local_block_ids = local_block_ids[-block_window:]
2340
- layer_remote_block_ids = remote_block_ids[-block_window:]
2341
-
2342
- # Get descs ids for the layer.
2343
- layer_local_desc_ids = self._get_block_descs_ids(
2344
- self.engine_id,
2345
- layer_local_block_ids,
2346
- layer_idx,
2347
- block_size_ratio=block_size_ratio,
2348
- )
2349
- layer_remote_desc_ids = self._get_block_descs_ids(
2350
- dst_engine_id,
2351
- layer_remote_block_ids,
2352
- layer_idx,
2353
- )
2354
-
2355
- local_descs_list.append(layer_local_desc_ids)
2356
- remote_descs_list.append(layer_remote_desc_ids)
2357
-
2358
- local_block_descs_ids = np.concatenate(local_descs_list)
2359
- remote_block_descs_ids = np.concatenate(remote_descs_list)
2412
+ remote_block_descs_ids = self._get_block_descs_ids(
2413
+ dst_engine_id,
2414
+ remote_block_ids,
2415
+ )
2416
+ local_block_descs_ids = self._get_block_descs_ids(
2417
+ self.engine_id,
2418
+ local_block_ids,
2419
+ block_size_ratio=block_size_ratio,
2420
+ )
2421
2422
assert len(local_block_descs_ids) == len(remote_block_descs_ids)
2423
@@ -2387,14 +2448,18 @@ class NixlConnectorWorker:
2448
dst_engine_id=dst_engine_id,
2449
remote_rank=remote_rank,
2450
)
2390
- if meta := self._recving_metadata.get(request_id):
2391
- self._invalid_block_ids.update(meta.local_block_ids)
2451
+ if (
2452
+ meta := self._recving_metadata.get(request_id)
2453
+ ) and not self._is_hma_required:
2454
+ self._invalid_block_ids.update(meta.local_block_ids[0])
2455
self.xfer_stats.record_failed_transfer()
2456
if handle is not None:
2457
self.nixl_wrapper.release_xfer_handle(handle)
2458
self._failed_recv_reqs.add(request_id)
2459
2397
- def get_mapped_blocks(self, block_ids, block_size_ratio):
2460
+ def get_mapped_blocks(
2461
+ self, block_ids: np.ndarray, block_size_ratio: int
2462
+ ) -> np.ndarray:
2463
"""
2464
Calculates the new set of block IDs by mapping every element
2465
in the (potentially sparse) input array.
@@ -2416,41 +2481,32 @@ class NixlConnectorWorker:
2481
def _get_block_descs_ids(
2482
self,
2483
engine_id: str,
2419
- block_ids: list[int],
2420
- layer_idx: int | None = None,
2484
+ block_ids: BlockIds,
2485
block_size_ratio: float | None = None,
2486
) -> np.ndarray:
2487
"""
2488
Get the descs ids for a set of block ids.
2425
- If layer_idx is provided, we use the region_ids for the given layer.
2426
- Otherwise, we use all regions.
2489
+ When HMA is enabled number of descriptors across kv cache groups might differ.
2490
+ A single flattened array is returned for all groups anyway.
2491
"""
2428
- if layer_idx is None:
2429
- region_ids = np.arange(self.num_regions)
2430
- else:
2431
- assert layer_idx < self.num_layers
2432
- if self.num_layers < self.num_regions:
2433
- # If we have more regions than layers, we assume that
2434
- # the regions are organized as [K0, V0, K1, V1, ...]
2435
- # and we select K_i and V_i
2436
- assert 2 * self.num_layers == self.num_regions
2437
- region_ids = np.arange(2 * layer_idx, 2 * layer_idx + 2)
2438
- else:
2439
- # Otherwise, we assume we have MLA and select i-th layer
2440
- assert self.num_layers == self.num_regions
2441
- region_ids = np.arange(layer_idx, layer_idx + 1)
2442
-
2492
+ region_ids = np.arange(self.num_regions)
2493
+ # NOTE (NickLucche) With HMA, every kv group has the same number of layers and
2494
+ # layers from different groups share the same kv tensor.
2495
+ # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions,
2496
+ # same for [3], but group0-group1 blocks will always differ (different areas).
2497
+ # Therefore we can just flatten the block_ids and compute the descs ids for all
2498
+ # groups at once.
2499
num_blocks = self.dst_num_blocks[engine_id]
2500
if block_size_ratio is not None:
2501
num_blocks = int(num_blocks * block_size_ratio)
2502
2503
# Compute the desc ids for each block.
2504
region_ids = region_ids[:, None]
2449
- block_ids = np.array(block_ids)[None, :]
2505
+ block_ids = np.concatenate(block_ids)[None, :]
2506
descs_ids = region_ids * num_blocks + block_ids
2507
return descs_ids.flatten()
2508
2453
- def _logical_to_kernel_block_ids(self, block_ids: list[int]) -> list[int]:
2509
+ def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
2510
"""
2511
Convert logical block ids to kernel physical block ids.
2512
This is required when the logical block size (the one set by the user)
@@ -2459,13 +2515,17 @@ class NixlConnectorWorker:
2515
if self._physical_blocks_per_logical_kv_block == 1:
2516
# Noop when physical and logical block sizes are the same
2517
return block_ids
2462
- block_ids_np = np.array(block_ids)
2518
block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
2519
1, -1
2520
)
2466
- return BlockTable.map_to_kernel_blocks(
2467
- block_ids_np, self._physical_blocks_per_logical_kv_block, block_arange
2468
- ).tolist()
2521
+ return [
2522
+ BlockTable.map_to_kernel_blocks(
2523
+ np.array(group),
2524
+ self._physical_blocks_per_logical_kv_block,
2525
+ block_arange,
2526
+ ).tolist()
2527
+ for group in block_ids
2528
+ ]
2529
2530
def get_backend_aware_kv_block_len(self, layer_idx: int) -> int:
2531
"""
vllm/v1/core/kv_cache_manager.py
+12
@@ -84,6 +84,18 @@ class KVCacheBlocks:
84
assert len(self.blocks) == 1, "Only one group is supported"
85
return [block.block_id for block in self.blocks[0] if block.block_hash is None]
86
87
+ def get_unhashed_block_ids_all_groups(self) -> list[list[int]]:
88
+ """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
89
+ # Skip padding blocks.
90
+ return [
91
+ [
92
+ block.block_id
93
+ for block in group
94
+ if block.block_hash is None and not block.is_null
95
+ ]
96
+ for group in self.blocks
97
+ ]
98
+
99
def new_empty(self) -> "KVCacheBlocks":
100
"""
101
Creates a new KVCacheBlocks instance with no blocks.