@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7ebc0ec104
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 11:00 UTC

[CI] Nixl+SimpleCPUOffloadingConnector unit tests (#43871)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed May 29, 2026 at 11:40 UTC 7ebc0ec1040604ebc586880ae9175af89e6cbcbe
1 file changed +325
tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py new
+325
@@ -0,0 +1,325 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for MultiConnector wrapping NixlConnector + SimpleCPUOffloadConnector.
4 +
5 +Verifies scheduler-level behavior: HMA support detection, load delegation
6 +(first-wins), store-to-all, and connector metadata aggregation.
7 +"""
8 +
9 +import uuid
10 +from collections import defaultdict
11 +from unittest.mock import patch
12 +
13 +import pytest
14 +import torch
15 +
16 +from tests.v1.kv_connector.unit.utils import (
17 + create_model_runner_output,
18 + create_request,
19 + create_scheduler,
20 + create_vllm_config,
21 +)
22 +from vllm import SamplingParams
23 +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
24 + MultiConnector,
25 + MultiKVConnectorMetadata,
26 +)
27 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
28 + NixlConnectorMetadata,
29 +)
30 +from vllm.utils.hashing import sha256
31 +from vllm.v1.core.kv_cache_utils import get_request_block_hasher
32 +from vllm.v1.kv_cache_interface import (
33 + FullAttentionSpec,
34 + KVCacheConfig,
35 + KVCacheGroupSpec,
36 + KVCacheTensor,
37 + SlidingWindowSpec,
38 +)
39 +from vllm.v1.outputs import KVConnectorOutput
40 +from vllm.v1.request import Request
41 +from vllm.v1.simple_kv_offload.metadata import (
42 + SimpleCPUOffloadMetadata,
43 + SimpleCPUOffloadWorkerMetadata,
44 +)
45 +
46 +NIXL_WRAPPER_PATCH = (
47 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"
48 +)
49 +
50 +
51 +class FakeNixlWrapper:
52 + """Minimal mock of NixlWrapper for testing without NIXL hardware.
53 +
54 + Duplicated from test_nixl_connector.py to avoid importing that module
55 + (which has heavy dependencies like ray).
56 + """
57 +
58 + AGENT_METADATA = b"fake_agent_metadata"
59 + REMOTE_AGENT_NAME = "remote_agent"
60 +
61 + def __init__(self, agent_name: str, *args, **kwargs):
62 + self._cycles_before_xfer_done = 0
63 + self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(lambda: 0)
64 +
65 + def get_reg_descs(self, caches_data, memory_type: str) -> list:
66 + return [str(uuid.uuid4()) for _ in caches_data]
67 +
68 + def register_memory(self, descs, backends) -> None:
69 + pass
70 +
71 + def deregister_memory(self, descs) -> None:
72 + pass
73 +
74 + def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
75 + return [str(uuid.uuid4()) for _ in blocks_data]
76 +
77 + def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
78 + return uuid.uuid4().int
79 +
80 + def get_agent_metadata(self) -> bytes:
81 + return self.AGENT_METADATA
82 +
83 + def add_remote_agent(self, agent_metadata: bytes) -> str:
84 + return self.REMOTE_AGENT_NAME
85 +
86 + def get_new_notifs(self) -> dict[str, list[bytes]]:
87 + return {}
88 +
89 + def check_xfer_state(self, handle: int) -> str:
90 + if self._check_xfer_state_cycles[handle] >= self._cycles_before_xfer_done:
91 + return "DONE"
92 + self._check_xfer_state_cycles[handle] += 1
93 + return "PROC"
94 +
95 + def release_xfer_handle(self, handle: int) -> None:
96 + pass
97 +
98 + def release_dlist_handle(self, handle: int) -> None:
99 + pass
100 +
101 + def remove_remote_agent(self, agent: str) -> None:
102 + pass
103 +
104 + def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
105 + pass
106 +
107 + def make_prepped_xfer(self, *args, **kwargs) -> int:
108 + return uuid.uuid4().int
109 +
110 + def transfer(self, handle: int) -> str:
111 + return "PROC"
112 +
113 + def get_xfer_telemetry(self, handle: int) -> dict:
114 + return {}
115 +
116 + def set_cycles_before_xfer_done(self, cycles: int):
117 + pass
118 +
119 +
120 +BLOCK_SIZE = 16
121 +NUM_KV_HEADS = 1
122 +HEAD_SIZE = 16
123 +DTYPE = torch.float16
124 +_BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize
125 +
126 +
127 +def _make_kv_cache_config(
128 + num_blocks: int = 100,
129 + swa_enabled: bool = False,
130 +) -> KVCacheConfig:
131 + """Build KVCacheConfig with non-empty kv_cache_tensors.
132 +
133 + SimpleCPUOffloadConnector requires kv_cache_tensors with real sizes
134 + (used by _derive_cpu_config).
135 + """
136 + fa_layers = ["layer0", "layer2"]
137 + groups = [
138 + KVCacheGroupSpec(
139 + fa_layers,
140 + FullAttentionSpec(
141 + block_size=BLOCK_SIZE,
142 + num_kv_heads=NUM_KV_HEADS,
143 + head_size=HEAD_SIZE,
144 + dtype=DTYPE,
145 + ),
146 + )
147 + ]
148 + tensors = [
149 + KVCacheTensor(
150 + size=_BYTES_PER_BLOCK * num_blocks,
151 + shared_by=fa_layers,
152 + )
153 + ]
154 + if swa_enabled:
155 + sw_layers = ["layer1", "layer3"]
156 + groups.append(
157 + KVCacheGroupSpec(
158 + sw_layers,
159 + SlidingWindowSpec(
160 + block_size=BLOCK_SIZE,
161 + num_kv_heads=NUM_KV_HEADS,
162 + head_size=HEAD_SIZE,
163 + dtype=DTYPE,
164 + sliding_window=128,
165 + ),
166 + )
167 + )
168 + tensors.append(
169 + KVCacheTensor(
170 + size=_BYTES_PER_BLOCK * num_blocks,
171 + shared_by=sw_layers,
172 + )
173 + )
174 + return KVCacheConfig(
175 + num_blocks=num_blocks,
176 + kv_cache_tensors=tensors,
177 + kv_cache_groups=groups,
178 + )
179 +
180 +
181 +def _multi_connector_config(swa_enabled: bool = False):
182 + """Return (vllm_config, kv_cache_config) for a MultiConnector test."""
183 + kv_cache_config = _make_kv_cache_config(swa_enabled=swa_enabled)
184 + vllm_config = create_vllm_config(
185 + kv_connector="MultiConnector",
186 + kv_connector_extra_config={
187 + "connectors": [
188 + {
189 + "kv_connector": "NixlConnector",
190 + "kv_role": "kv_both",
191 + },
192 + {
193 + "kv_connector": "SimpleCPUOffloadConnector",
194 + "kv_role": "kv_both",
195 + "kv_connector_extra_config": {
196 + "cpu_bytes_to_use": 1 << 30,
197 + },
198 + },
199 + ],
200 + },
201 + )
202 + return vllm_config, kv_cache_config
203 +
204 +
205 +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
206 +def test_nixl_wins_load_over_cpu_offload():
207 + """When NixlConnector (index 0) has matched tokens from a remote prefill, it should
208 + win the load: Nixl metadata tracks the recv while CPU offload metadata has no load
209 + scheduled."""
210 + vllm_config, kv_cache_config = _multi_connector_config()
211 + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
212 + mc = scheduler.connector
213 + assert isinstance(mc, MultiConnector)
214 +
215 + request = create_request(
216 + request_id=1,
217 + num_tokens=BLOCK_SIZE * 3,
218 + do_remote_prefill=True,
219 + block_size=BLOCK_SIZE,
220 + )
221 + scheduler.add_request(request)
222 + sched_out = scheduler.schedule()
223 +
224 + assert mc._requests_to_connector[request.request_id] == 0
225 +
226 + meta = sched_out.kv_connector_metadata
227 + assert isinstance(meta, MultiKVConnectorMetadata)
228 + assert len(meta.metadata) == 2
229 +
230 + nixl_meta = meta.metadata[0]
231 + assert isinstance(nixl_meta, NixlConnectorMetadata)
232 + # nixl is tracking the request
233 + assert request.request_id in nixl_meta.reqs_to_recv
234 +
235 + cpu_meta = meta.metadata[1]
236 + assert isinstance(cpu_meta, SimpleCPUOffloadMetadata)
237 + assert not cpu_meta.load_gpu_blocks
238 +
239 +
240 +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
241 +def test_cpu_offload_wins_when_nixl_has_no_match():
242 + """When NixlConnector returns 0 matched tokens and SimpleCPUOffloadConnector has a
243 + CPU cache hit, the CPU offload connector (index 1) wins the load."""
244 + vllm_config, kv_cache_config = _multi_connector_config()
245 + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
246 + mc = scheduler.connector
247 + assert isinstance(mc, MultiConnector)
248 +
249 + req1 = create_request(
250 + request_id=10,
251 + num_tokens=BLOCK_SIZE * 3,
252 + block_size=BLOCK_SIZE,
253 + )
254 + scheduler.add_request(req1)
255 + sched_out1 = scheduler.schedule()
256 +
257 + # build_connector_meta runs before _update_after_schedule, so num_computed_tokens
258 + # is still 0 during the first schedule and the store sees no confirmed blocks.
259 + # Simulate one model step so the second schedule triggers the store.
260 + model_output = create_model_runner_output(reqs=[req1])
261 + scheduler.update_from_output(sched_out1, model_output)
262 +
263 + sched_out2 = scheduler.schedule()
264 + meta2 = sched_out2.kv_connector_metadata
265 + assert isinstance(meta2, MultiKVConnectorMetadata)
266 + cpu_meta = meta2.metadata[1]
267 + assert isinstance(cpu_meta, SimpleCPUOffloadMetadata)
268 + assert cpu_meta.store_event >= 0, "Expected a store event on the second schedule"
269 +
270 + cpu_connector = mc._connectors[1]
271 + worker_meta = SimpleCPUOffloadWorkerMetadata(
272 + completed_store_events={
273 + cpu_meta.store_event: cpu_connector.scheduler_manager._expected_worker_count
274 + },
275 + )
276 + output = KVConnectorOutput(
277 + finished_recving=set(),
278 + kv_connector_worker_meta=worker_meta,
279 + )
280 + cpu_connector.update_connector_output(output)
281 +
282 + req2 = Request(
283 + request_id="id-cpu-offload-hit",
284 + prompt_token_ids=req1.prompt_token_ids,
285 + sampling_params=SamplingParams(max_tokens=16),
286 + pooling_params=None,
287 + mm_features=None,
288 + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
289 + )
290 +
291 + hit_tokens, is_async = mc.get_num_new_matched_tokens(req2, num_computed_tokens=0)
292 + assert hit_tokens is not None and hit_tokens > 0
293 + assert mc._requests_to_connector[req2.request_id] == 1
294 + assert is_async is True
295 +
296 +
297 +@pytest.mark.parametrize("swa_enabled", [False, True], ids=["fa_only", "fa_sw"])
298 +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper)
299 +def test_request_finished_no_async_save(swa_enabled: bool):
300 + """A normal request (no P/D) produces no async save from either connector.
301 + MultiConnector returns (False, None) via both request_finished and
302 + request_finished_all_groups, and cleans up _requests_to_connector."""
303 + from vllm.v1.request import RequestStatus
304 +
305 + vllm_config, kv_cache_config = _multi_connector_config(swa_enabled=swa_enabled)
306 + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)
307 + mc = scheduler.connector
308 + assert isinstance(mc, MultiConnector)
309 +
310 + request = create_request(
311 + request_id=40,
312 + num_tokens=BLOCK_SIZE * 2,
313 + block_size=BLOCK_SIZE,
314 + )
315 + scheduler.add_request(request)
316 + scheduler.schedule()
317 + request.status = RequestStatus.FINISHED_STOPPED
318 +
319 + block_ids = (list(range(2)), list(range(2))) if swa_enabled else (list(range(2)),)
320 + async_save, txfer_params = mc.request_finished_all_groups(request, block_ids)
321 +
322 + assert async_save is False
323 + assert txfer_params is None
324 + assert len(mc._extra_async_saves) == 0
325 + assert request.request_id not in mc._requests_to_connector