| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import socket |
| 4 | import time |
| 5 | |
| 6 | import msgspec |
| 7 | import msgspec.msgpack |
| 8 | import prometheus_client |
| 9 | import pytest |
| 10 | import zmq |
| 11 | from tqdm import tqdm |
| 12 | |
| 13 | from vllm import LLM, SamplingParams, TokensPrompt |
| 14 | from vllm.config import KVEventsConfig, KVTransferConfig |
| 15 | from vllm.distributed.kv_events import BlockStored, KVEventBatch |
| 16 | from vllm.platforms import current_platform |
| 17 | |
| 18 | # The offloaded block size must stay a whole number of GPU blocks. |
| 19 | CPU_BLOCK_SIZES: int |
| 20 | if current_platform.is_xpu(): |
| 21 | CPU_BLOCK_SIZES = 64 |
| 22 | elif current_platform.is_rocm(): |
| 23 | # ROCM_AITER_UNIFIED_ATTN prefers 64-token GPU blocks, so 192 covers both that |
| 24 | # and the 16-token default. |
| 25 | CPU_BLOCK_SIZES = 192 |
| 26 | else: |
| 27 | CPU_BLOCK_SIZES = 48 |
| 28 | |
| 29 | _ATTN_BACKENDS: list[str] = [] |
| 30 | if current_platform.is_cuda(): |
| 31 | _ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"] |
| 32 | elif current_platform.is_rocm(): |
| 33 | from vllm._aiter_ops import is_aiter_found_and_supported |
| 34 | |
| 35 | _ATTN_BACKENDS = ["TRITON_ATTN"] |
| 36 | if is_aiter_found_and_supported(): |
| 37 | _ATTN_BACKENDS.append("ROCM_AITER_UNIFIED_ATTN") |
| 38 | elif current_platform.is_xpu(): |
| 39 | _ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"] |
| 40 | |
| 41 | # (model, attn_backend | None, block_size | None, uses_hma) |
| 42 | # |
| 43 | # - Llama: tested with each attention backend and a custom block_size. |
| 44 | # - Gemma-3: HMA (sliding window + full attention), default backend. |
| 45 | # - Mamba-130m: HMA (attention-free, acts like sliding_window=1), |
| 46 | # default backend. Prefix caching must be force-enabled. |
| 47 | # - Falcon-H1-0.5B-Instruct: HMA (parallel SSM/attention in every layer). |
| 48 | # After page-size unification the mamba and attention groups have |
| 49 | # different block sizes. |
| 50 | MODEL_PARAMS: list[tuple[str, str | None, int | None, bool]] = [ |
| 51 | ("meta-llama/Llama-3.2-1B-Instruct", backend, CPU_BLOCK_SIZES, False) |
| 52 | for backend in _ATTN_BACKENDS |
| 53 | ] |
| 54 | # HMA / Mamba models are only tested on CUDA (not ROCm). |
| 55 | if current_platform.is_cuda(): |
| 56 | MODEL_PARAMS += [ |
| 57 | ("google/gemma-3-1b-it", None, CPU_BLOCK_SIZES, True), |
| 58 | ("state-spaces/mamba-130m-hf", None, CPU_BLOCK_SIZES, True), |
| 59 | # Falcon-H1: parallel hybrid (every layer has both attention and SSM). |
| 60 | # The mamba and attention groups end up with different GPU block sizes |
| 61 | # after page-size unification, so we leave cpu_block_size=None |
| 62 | # (blocks_per_chunk stays 1). |
| 63 | ("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True), |
| 64 | ] |
| 65 | |
| 66 | # Maximum time (seconds) to wait for the async CPU offload transfer |
| 67 | # to complete before giving up. |
| 68 | _RESET_CACHE_TIMEOUT = 30 if current_platform.is_rocm() else 10 |
| 69 | |
| 70 | # ZMQ poll timeout (ms) for the first event. |
| 71 | _FIRST_EVENT_POLL_MS = 10_000 if current_platform.is_rocm() else 1000 |
| 72 | |
| 73 | # Hard ceiling (seconds) on how long get_new_cpu_stored_events may loop, |
| 74 | # to prevent hangs if non-CPU events keep arriving indefinitely. |
| 75 | _EVENT_DRAIN_TIMEOUT = 60 |
| 76 | |
| 77 | |
| 78 | class MockSubscriber: |
| 79 | """Helper class to receive and verify published events.""" |
| 80 | |
| 81 | def __init__( |
| 82 | self, |
| 83 | endpoint: str, |
| 84 | topic: str, |
| 85 | ): |
| 86 | self.ctx = zmq.Context.instance() |
| 87 | self.topic_bytes = topic.encode("utf-8") |
| 88 | |
| 89 | # Set up subscriber socket |
| 90 | self.sub = self.ctx.socket(zmq.SUB) |
| 91 | self.sub.setsockopt(zmq.SUBSCRIBE, self.topic_bytes) |
| 92 | self.sub.connect(endpoint) |
| 93 | |
| 94 | self.decoder = msgspec.msgpack.Decoder(type=KVEventBatch) |
| 95 | |
| 96 | def get_new_cpu_stored_events(self) -> list[BlockStored]: |
| 97 | cpu_stored_events: list[BlockStored] = [] |
| 98 | |
| 99 | poller = zmq.Poller() |
| 100 | poller.register(self.sub, zmq.POLLIN) |
| 101 | |
| 102 | poll_ms = _FIRST_EVENT_POLL_MS |
| 103 | deadline = time.monotonic() + _EVENT_DRAIN_TIMEOUT |
| 104 | while time.monotonic() < deadline: |
| 105 | events = dict(poller.poll(poll_ms)) |
| 106 | |
| 107 | if events.get(self.sub) != zmq.POLLIN: |
| 108 | return cpu_stored_events |
| 109 | |
| 110 | topic_bytes, _, payload = self.sub.recv_multipart() |
| 111 | |
| 112 | assert topic_bytes == self.topic_bytes |
| 113 | |
| 114 | event_batch = self.decoder.decode(payload) |
| 115 | assert isinstance(event_batch, KVEventBatch) |
| 116 | for event in event_batch.events: |
| 117 | if isinstance(event, BlockStored) and event.medium == "CPU": |
| 118 | cpu_stored_events.append(event) |
| 119 | poll_ms = 100 |
| 120 | |
| 121 | return cpu_stored_events |
| 122 | |
| 123 | def close(self): |
| 124 | """Clean up resources.""" |
| 125 | self.sub.close() |
| 126 | |
| 127 | |
| 128 | def _wait_for_prefix_cache_reset(llm: LLM) -> None: |
| 129 | """Wait for async offload transfers to finish so prefix cache can reset. |
| 130 | |
| 131 | The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks |
| 132 | are still held by the offload worker, ``reset_prefix_cache`` returns |
| 133 | ``False``. Between retries we send a dummy single-token prefill to force |
| 134 | the engine to step, which polls the worker for completed transfers and |
| 135 | frees GPU blocks. |
| 136 | """ |
| 137 | _dummy_params = SamplingParams(max_tokens=1) |
| 138 | deadline = time.monotonic() + _RESET_CACHE_TIMEOUT |
| 139 | while not llm.reset_prefix_cache(): |
| 140 | if time.monotonic() > deadline: |
| 141 | raise TimeoutError( |
| 142 | "reset_prefix_cache did not succeed within " |
| 143 | f"{_RESET_CACHE_TIMEOUT}s - async offload may be stuck" |
| 144 | ) |
| 145 | # Force an engine step so the scheduler polls get_finished() |
| 146 | # and releases GPU blocks held by in-flight async stores. |
| 147 | llm.generate( |
| 148 | [TokensPrompt(prompt_token_ids=[0])], |
| 149 | _dummy_params, |
| 150 | use_tqdm=False, |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def _latency_test(llm: LLM, subscriber: MockSubscriber | None): |
| 155 | # TODO: Reintroduce latency test on ROCm once MRV2 supports cross |
| 156 | # layer KV Cache. See https://github.com/vllm-project/vllm/pull/45947 |
| 157 | if current_platform.is_rocm(): |
| 158 | return |
| 159 | sampling_params = SamplingParams(max_tokens=1) |
| 160 | |
| 161 | num_times_cpu_better_than_cold = 0 |
| 162 | num_tests = 10 |
| 163 | total_cold_time = 0.0 |
| 164 | total_gpu_hit_time = 0.0 |
| 165 | total_cpu_hit_time = 0.0 |
| 166 | max_model_len = llm.llm_engine.vllm_config.model_config.max_model_len |
| 167 | # Use a long prompt that fits within the model's context window. |
| 168 | prompt_len = min(10001, max_model_len - 1) |
| 169 | prompt_token_ids = [0] * prompt_len |
| 170 | for i in tqdm(range(num_tests), desc="Running tests"): |
| 171 | prompt_token_ids[0] = i |
| 172 | prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)] |
| 173 | |
| 174 | # run generation - this should trigger saving KV cache |
| 175 | start_time = time.time() |
| 176 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 177 | cold_time = time.time() - start_time |
| 178 | total_cold_time += cold_time |
| 179 | |
| 180 | # run generation again - should hit the GPU prefix cache |
| 181 | start_time = time.time() |
| 182 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 183 | gpu_hit_time = time.time() - start_time |
| 184 | total_gpu_hit_time += gpu_hit_time |
| 185 | |
| 186 | # Wait for the async CPU offload to finish, then reset prefix cache |
| 187 | # so the next generate() must reload from CPU rather than GPU. |
| 188 | _wait_for_prefix_cache_reset(llm) |
| 189 | |
| 190 | # Verify CPU stored events arrived (offload is done before we |
| 191 | # attempt to load from CPU). |
| 192 | if subscriber is not None: |
| 193 | assert subscriber.get_new_cpu_stored_events(), ( |
| 194 | f"No CPU stored events received on iteration {i}; " |
| 195 | "async offload may not have completed in time" |
| 196 | ) |
| 197 | |
| 198 | # run generation again - this should trigger loading from CPU |
| 199 | start_time = time.time() |
| 200 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 201 | cpu_hit_time = time.time() - start_time |
| 202 | total_cpu_hit_time += cpu_hit_time |
| 203 | |
| 204 | if cpu_hit_time < cold_time: |
| 205 | num_times_cpu_better_than_cold += 1 |
| 206 | |
| 207 | print("Average times:") |
| 208 | print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms") |
| 209 | print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms") |
| 210 | print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms") |
| 211 | |
| 212 | assert num_times_cpu_better_than_cold >= 0.8 * num_tests |
| 213 | |
| 214 | |
| 215 | def _accuracy_test(llm: LLM, subscriber: MockSubscriber | None): |
| 216 | sampling_params = SamplingParams(max_tokens=1) |
| 217 | extra_config = ( |
| 218 | llm.llm_engine.vllm_config.kv_transfer_config.kv_connector_extra_config |
| 219 | ) |
| 220 | cpu_block_size = extra_config.get("block_size") |
| 221 | if cpu_block_size is None: |
| 222 | # No custom offloaded block_size: offloaded blocks match GPU blocks. |
| 223 | # Use the hash block_size (cache_config.block_size) for alignment. |
| 224 | cpu_block_size = llm.llm_engine.vllm_config.cache_config.block_size |
| 225 | |
| 226 | if subscriber is not None: |
| 227 | subscriber.get_new_cpu_stored_events() |
| 228 | |
| 229 | # Align the reusable prefix; the final token must be recomputed for logits. |
| 230 | tokenizer = llm.get_tokenizer() |
| 231 | prompt_token_ids = tokenizer.encode("Let's count to 10. One, two, three, four,") |
| 232 | reusable_prefix = prompt_token_ids[:-1] |
| 233 | final_token = prompt_token_ids[-1:] |
| 234 | |
| 235 | [padding_token_id] = tokenizer.encode(".", add_special_tokens=False) |
| 236 | padding = [padding_token_id] * (-len(reusable_prefix) % cpu_block_size) |
| 237 | insert_at = int(reusable_prefix[0] == tokenizer.bos_token_id) |
| 238 | reusable_prefix[insert_at:insert_at] = padding |
| 239 | assert len(reusable_prefix) % cpu_block_size == 0 |
| 240 | |
| 241 | prompt = TokensPrompt( |
| 242 | prompt_token_ids=reusable_prefix + final_token, |
| 243 | ) |
| 244 | |
| 245 | # Seed the CPU cache with the prompt. |
| 246 | llm.generate(prompt, sampling_params, use_tqdm=False) |
| 247 | |
| 248 | if subscriber is not None: |
| 249 | assert subscriber.get_new_cpu_stored_events() |
| 250 | |
| 251 | test_count = 20 |
| 252 | results = llm.generate([prompt] * test_count, sampling_params, use_tqdm=False) |
| 253 | success_count = sum(1 for r in results if r.outputs[0].text == " five") |
| 254 | assert success_count >= 0.5 * test_count |
| 255 | |
| 256 | |
| 257 | @pytest.mark.parametrize("model, attn_backend, cpu_block_size, uses_hma", MODEL_PARAMS) |
| 258 | def test_cpu_offloading( |
| 259 | model: str, |
| 260 | attn_backend: str | None, |
| 261 | cpu_block_size: int | None, |
| 262 | uses_hma: bool, |
| 263 | ) -> None: |
| 264 | """Tests OffloadingConnector with CPUOffloadingSpec.""" |
| 265 | # configure OffloadingConnector (spec_name=CPUOffloadingSpec by default) |
| 266 | extra_config: dict = {"cpu_bytes_to_use": 500 << 20} |
| 267 | if cpu_block_size is not None: |
| 268 | extra_config["block_size"] = cpu_block_size |
| 269 | kv_transfer_config = KVTransferConfig( |
| 270 | kv_connector="OffloadingConnector", |
| 271 | kv_role="kv_both", |
| 272 | kv_connector_extra_config=extra_config, |
| 273 | ) |
| 274 | |
| 275 | subscriber: MockSubscriber | None = None |
| 276 | port: int |
| 277 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 278 | s.bind(("0.0.0.0", 0)) |
| 279 | port = s.getsockname()[1] |
| 280 | |
| 281 | events_endpoint = f"tcp://*:{port}" |
| 282 | kv_events_config = KVEventsConfig( |
| 283 | enable_kv_cache_events=True, |
| 284 | publisher="zmq", |
| 285 | endpoint=events_endpoint, |
| 286 | topic="test", |
| 287 | ) |
| 288 | |
| 289 | # Attention-free / hybrid models disable prefix caching by default |
| 290 | # (ModelConfig.is_prefix_caching_supported returns False). Without it, |
| 291 | # mamba_block_size falls back to max_model_len, making GPU blocks too |
| 292 | # large for any reasonable offloaded block_size. Force-enable it. |
| 293 | force_prefix_caching = uses_hma |
| 294 | |
| 295 | llm = LLM( |
| 296 | model=model, |
| 297 | max_model_len=4096, |
| 298 | gpu_memory_utilization=0.5, |
| 299 | kv_events_config=kv_events_config, |
| 300 | kv_transfer_config=kv_transfer_config, |
| 301 | **({"attention_config": {"backend": attn_backend}} if attn_backend else {}), |
| 302 | # Keep HMA explicitly enabled for HMA model coverage. |
| 303 | **({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}), |
| 304 | **({"enable_prefix_caching": True} if force_prefix_caching else {}), |
| 305 | # ROCm: batch size 1 to reduce variability |
| 306 | **({"max_num_seqs": 1} if current_platform.is_rocm() else {}), |
| 307 | ) |
| 308 | |
| 309 | events_endpoint = events_endpoint.replace("*", "127.0.0.1") |
| 310 | subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic) |
| 311 | |
| 312 | try: |
| 313 | _latency_test(llm, subscriber) |
| 314 | _accuracy_test(llm, subscriber) |
| 315 | finally: |
| 316 | if subscriber is not None: |
| 317 | subscriber.close() |
| 318 | del llm |
| 319 | |
| 320 | |
| 321 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") |
| 322 | def test_cpu_offloading_metrics() -> None: |
| 323 | """Verify that offloading Prometheus metrics (new flat and deprecated |
| 324 | labeled) are emitted after stores and loads.""" |
| 325 | extra_config: dict = { |
| 326 | "cpu_bytes_to_use": 500 << 20, |
| 327 | "block_size": CPU_BLOCK_SIZES, |
| 328 | } |
| 329 | kv_transfer_config = KVTransferConfig( |
| 330 | kv_connector="OffloadingConnector", |
| 331 | kv_role="kv_both", |
| 332 | kv_connector_extra_config=extra_config, |
| 333 | ) |
| 334 | |
| 335 | llm = LLM( |
| 336 | model="meta-llama/Llama-3.2-1B-Instruct", |
| 337 | max_model_len=4096, |
| 338 | gpu_memory_utilization=0.5, |
| 339 | kv_transfer_config=kv_transfer_config, |
| 340 | disable_log_stats=False, |
| 341 | ) |
| 342 | |
| 343 | try: |
| 344 | prompt_token_ids = list(range(500)) |
| 345 | |
| 346 | # First generate: cold run, triggers a store to CPU. |
| 347 | # Use max_tokens>1 so the request is still producing output |
| 348 | # tokens when the async store completes and stats get drained. |
| 349 | # (The LLMEngine only records stats on steps with request outputs.) |
| 350 | llm.generate( |
| 351 | [TokensPrompt(prompt_token_ids=prompt_token_ids)], |
| 352 | SamplingParams(max_tokens=10), |
| 353 | use_tqdm=False, |
| 354 | ) |
| 355 | |
| 356 | # Wait for the async offload to finish, then reset GPU prefix cache |
| 357 | # so the next generate must load from CPU. |
| 358 | _wait_for_prefix_cache_reset(llm) |
| 359 | |
| 360 | # Second generate: triggers a load from CPU. |
| 361 | # Send a short filler alongside the load prompt so there's always |
| 362 | # a request producing output tokens when the load stats get drained |
| 363 | # (the LLMEngine only records stats on steps with request outputs). |
| 364 | filler = TokensPrompt(prompt_token_ids=[0]) |
| 365 | llm.generate( |
| 366 | [filler, TokensPrompt(prompt_token_ids=prompt_token_ids)], |
| 367 | SamplingParams(max_tokens=50), |
| 368 | use_tqdm=False, |
| 369 | ) |
| 370 | |
| 371 | # Metric helpers. |
| 372 | registry = prometheus_client.REGISTRY |
| 373 | |
| 374 | def _get_counter_value( |
| 375 | name: str, labels: dict[str, str] | None = None |
| 376 | ) -> float: |
| 377 | total = 0.0 |
| 378 | for metric in registry.collect(): |
| 379 | if metric.name == name: |
| 380 | for sample in metric.samples: |
| 381 | if sample.name != name + "_total": |
| 382 | continue |
| 383 | if labels and not all( |
| 384 | sample.labels.get(k) == v for k, v in labels.items() |
| 385 | ): |
| 386 | continue |
| 387 | total += sample.value |
| 388 | return total |
| 389 | |
| 390 | def _get_histogram_count( |
| 391 | name: str, labels: dict[str, str] | None = None |
| 392 | ) -> float: |
| 393 | total = 0.0 |
| 394 | for metric in registry.collect(): |
| 395 | if metric.name == name: |
| 396 | for sample in metric.samples: |
| 397 | if sample.name != name + "_count": |
| 398 | continue |
| 399 | if labels and not all( |
| 400 | sample.labels.get(k) == v for k, v in labels.items() |
| 401 | ): |
| 402 | continue |
| 403 | total += sample.value |
| 404 | return total |
| 405 | |
| 406 | # New flat counter metrics |
| 407 | store_bytes = _get_counter_value("vllm:kv_offload_store_bytes") |
| 408 | assert store_bytes > 0, f"Expected store_bytes > 0, got {store_bytes}" |
| 409 | load_bytes = _get_counter_value("vllm:kv_offload_load_bytes") |
| 410 | assert load_bytes > 0, f"Expected load_bytes > 0, got {load_bytes}" |
| 411 | store_time = _get_counter_value("vllm:kv_offload_store_time") |
| 412 | assert store_time > 0, f"Expected store_time > 0, got {store_time}" |
| 413 | load_time = _get_counter_value("vllm:kv_offload_load_time") |
| 414 | assert load_time > 0, f"Expected load_time > 0, got {load_time}" |
| 415 | |
| 416 | # New flat histogram metrics |
| 417 | store_size_count = _get_histogram_count("vllm:kv_offload_store_size") |
| 418 | assert store_size_count > 0, ( |
| 419 | f"Expected store_size histogram observations > 0, got {store_size_count}" |
| 420 | ) |
| 421 | load_size_count = _get_histogram_count("vllm:kv_offload_load_size") |
| 422 | assert load_size_count > 0, ( |
| 423 | f"Expected load_size histogram observations > 0, got {load_size_count}" |
| 424 | ) |
| 425 | |
| 426 | # Deprecated labeled metrics — verify per transfer_type label. |
| 427 | load_label = {"transfer_type": "CPU_to_GPU"} |
| 428 | store_label = {"transfer_type": "GPU_to_CPU"} |
| 429 | |
| 430 | dep_load_bytes = _get_counter_value("vllm:kv_offload_total_bytes", load_label) |
| 431 | assert dep_load_bytes > 0, ( |
| 432 | f"Expected deprecated load bytes > 0, got {dep_load_bytes}" |
| 433 | ) |
| 434 | dep_store_bytes = _get_counter_value("vllm:kv_offload_total_bytes", store_label) |
| 435 | assert dep_store_bytes > 0, ( |
| 436 | f"Expected deprecated store bytes > 0, got {dep_store_bytes}" |
| 437 | ) |
| 438 | dep_load_time = _get_counter_value("vllm:kv_offload_total_time", load_label) |
| 439 | assert dep_load_time > 0, ( |
| 440 | f"Expected deprecated load time > 0, got {dep_load_time}" |
| 441 | ) |
| 442 | dep_store_time = _get_counter_value("vllm:kv_offload_total_time", store_label) |
| 443 | assert dep_store_time > 0, ( |
| 444 | f"Expected deprecated store time > 0, got {dep_store_time}" |
| 445 | ) |
| 446 | dep_load_size = _get_histogram_count("vllm:kv_offload_size", load_label) |
| 447 | assert dep_load_size > 0, ( |
| 448 | f"Expected deprecated load size observations > 0, got {dep_load_size}" |
| 449 | ) |
| 450 | dep_store_size = _get_histogram_count("vllm:kv_offload_size", store_label) |
| 451 | assert dep_store_size > 0, ( |
| 452 | f"Expected deprecated store size observations > 0, got {dep_store_size}" |
| 453 | ) |
| 454 | |
| 455 | # Flat and deprecated metrics must be consistent (dual-write). |
| 456 | assert store_bytes == dep_store_bytes |
| 457 | assert load_bytes == dep_load_bytes |
| 458 | assert store_time == dep_store_time |
| 459 | assert load_time == dep_load_time |
| 460 | assert store_size_count == dep_store_size |
| 461 | assert load_size_count == dep_load_size |
| 462 | finally: |
| 463 | del llm |
| 464 | |
| 465 | |
| 466 | def test_tiering_offloading() -> None: |
| 467 | """Tests OffloadingConnector with TieringOffloadingSpec.""" |
| 468 | extra_config: dict = { |
| 469 | "cpu_bytes_to_use": 500 << 20, |
| 470 | "block_size": CPU_BLOCK_SIZES, |
| 471 | "spec_name": "TieringOffloadingSpec", |
| 472 | "secondary_tiers": [{"type": "example"}], |
| 473 | } |
| 474 | kv_transfer_config = KVTransferConfig( |
| 475 | kv_connector="OffloadingConnector", |
| 476 | kv_role="kv_both", |
| 477 | kv_connector_extra_config=extra_config, |
| 478 | ) |
| 479 | |
| 480 | port: int |
| 481 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 482 | s.bind(("0.0.0.0", 0)) |
| 483 | port = s.getsockname()[1] |
| 484 | events_endpoint = f"tcp://*:{port}" |
| 485 | kv_events_config = KVEventsConfig( |
| 486 | enable_kv_cache_events=True, |
| 487 | publisher="zmq", |
| 488 | endpoint=events_endpoint, |
| 489 | topic="test", |
| 490 | ) |
| 491 | |
| 492 | llm = LLM( |
| 493 | model="meta-llama/Llama-3.2-1B-Instruct", |
| 494 | max_model_len=4096, |
| 495 | gpu_memory_utilization=0.5, |
| 496 | kv_events_config=kv_events_config, |
| 497 | kv_transfer_config=kv_transfer_config, |
| 498 | ) |
| 499 | subscriber = MockSubscriber( |
| 500 | events_endpoint.replace("*", "127.0.0.1"), |
| 501 | topic=kv_events_config.topic, |
| 502 | ) |
| 503 | try: |
| 504 | _latency_test(llm, subscriber) |
| 505 | _accuracy_test(llm, subscriber) |
| 506 | finally: |
| 507 | subscriber.close() |
| 508 | del llm |
| 509 | |
| 510 | |
| 511 | def test_fs_tiering_offloading(tmp_path) -> None: |
| 512 | """Tests OffloadingConnector with TieringOffloadingSpec |
| 513 | + fs secondary tier.""" |
| 514 | extra_config: dict = { |
| 515 | "cpu_bytes_to_use": 1 << 30, |
| 516 | "block_size": CPU_BLOCK_SIZES, |
| 517 | "spec_name": "TieringOffloadingSpec", |
| 518 | "secondary_tiers": [{"type": "fs", "root_dir": str(tmp_path)}], |
| 519 | } |
| 520 | kv_transfer_config = KVTransferConfig( |
| 521 | kv_connector="OffloadingConnector", |
| 522 | kv_role="kv_both", |
| 523 | kv_connector_extra_config=extra_config, |
| 524 | ) |
| 525 | |
| 526 | port: int |
| 527 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 528 | s.bind(("0.0.0.0", 0)) |
| 529 | port = s.getsockname()[1] |
| 530 | events_endpoint = f"tcp://*:{port}" |
| 531 | kv_events_config = KVEventsConfig( |
| 532 | enable_kv_cache_events=True, |
| 533 | publisher="zmq", |
| 534 | endpoint=events_endpoint, |
| 535 | topic="test", |
| 536 | ) |
| 537 | |
| 538 | llm = LLM( |
| 539 | model="meta-llama/Llama-3.2-1B-Instruct", |
| 540 | max_model_len=4096, |
| 541 | gpu_memory_utilization=0.5, |
| 542 | kv_events_config=kv_events_config, |
| 543 | kv_transfer_config=kv_transfer_config, |
| 544 | ) |
| 545 | subscriber = MockSubscriber( |
| 546 | events_endpoint.replace("*", "127.0.0.1"), |
| 547 | topic=kv_events_config.topic, |
| 548 | ) |
| 549 | try: |
| 550 | _latency_test(llm, subscriber) |
| 551 | _accuracy_test(llm, subscriber) |
| 552 | finally: |
| 553 | subscriber.close() |
| 554 | del llm |
| 555 | |
| 556 | |
| 557 | @pytest.mark.skipif( |
| 558 | not current_platform.is_cuda(), |
| 559 | reason="HMA mamba-align CPU offload test is CUDA-only", |
| 560 | ) |
| 561 | @pytest.mark.parametrize( |
| 562 | "model,block_size,tp_size", |
| 563 | [ |
| 564 | # ("Qwen/Qwen3.6-35B-A3B", 1056, 2), |
| 565 | # ("tiiuae/falcon-mamba-7b", 16, 1), |
| 566 | ("state-spaces/mamba-1.4b-hf", 16, 1), |
| 567 | ], |
| 568 | ) |
| 569 | def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int): |
| 570 | kv_transfer_config = KVTransferConfig( |
| 571 | kv_connector="OffloadingConnector", |
| 572 | kv_role="kv_both", |
| 573 | kv_connector_extra_config={ |
| 574 | "cpu_bytes_to_use": 4 << 30, |
| 575 | "block_size": block_size, |
| 576 | }, |
| 577 | ) |
| 578 | llm = LLM( |
| 579 | model=model, |
| 580 | max_model_len=block_size * 10, |
| 581 | gpu_memory_utilization=0.85, |
| 582 | tensor_parallel_size=tp_size, |
| 583 | kv_transfer_config=kv_transfer_config, |
| 584 | language_model_only=True, |
| 585 | enable_prefix_caching=True, |
| 586 | mamba_cache_mode="align", |
| 587 | # Use lossless state storage so this exact-equality regression isolates |
| 588 | # offload boundary selection from low-precision Mamba checkpointing. |
| 589 | mamba_cache_dtype="float32", |
| 590 | mamba_ssm_cache_dtype="float32", |
| 591 | disable_hybrid_kv_cache_manager=False, |
| 592 | ) |
| 593 | |
| 594 | _PROMPT_SIZE: int = block_size * 2 |
| 595 | # Cold prefill and offload-resume are not bit-exact (selective_scan_fn over |
| 596 | # the whole prompt vs selective_state_update for the recomputed token), so |
| 597 | # exact-text equality only holds while the rounding difference never flips a |
| 598 | # greedy argmax. Use a long, highly predictable prompt: it needs no "...." |
| 599 | # padding and keeps top-1 ahead of top-2 by >2.8 logprobs at every step. |
| 600 | _PROMPT_TEXT = ( |
| 601 | "The cat sat on the mat. The cat sat on the mat. The cat sat on the mat. " |
| 602 | "The cat sat on the mat. The cat sat on the mat. The cat sat on the mat. " |
| 603 | ) |
| 604 | |
| 605 | # build prompt ids to match prompt_size |
| 606 | tokenizer = llm.get_tokenizer() |
| 607 | raw_ids: list[int] = tokenizer.encode(_PROMPT_TEXT) |
| 608 | while len(raw_ids) < _PROMPT_SIZE: |
| 609 | raw_ids = tokenizer.encode("....") + raw_ids |
| 610 | initial_ids: list[int] = raw_ids[:_PROMPT_SIZE] |
| 611 | |
| 612 | sampling_params = SamplingParams(max_tokens=128, temperature=0, ignore_eos=True) |
| 613 | |
| 614 | failures: list[str] = [] |
| 615 | |
| 616 | def _get_output_str(outputs): |
| 617 | return outputs[0].outputs[0].text |
| 618 | |
| 619 | def _verify(llm, prompt, label: str): |
| 620 | cold_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) |
| 621 | _wait_for_prefix_cache_reset(llm) |
| 622 | cpu_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) |
| 623 | |
| 624 | cold_text = _get_output_str(cold_outputs) |
| 625 | cpu_text = _get_output_str(cpu_outputs) |
| 626 | print(f"{label} : cold outputs\n{cold_text}") |
| 627 | print(f"{label} : cpu outputs\n{cpu_text}") |
| 628 | |
| 629 | if cold_text != cpu_text: |
| 630 | failures.append( |
| 631 | f"{label}: mismatch\n cold: {cold_text!r}\n cpu: {cpu_text!r}" |
| 632 | ) |
| 633 | |
| 634 | try: |
| 635 | # A cached Mamba entry is a point state. CPU cache stores are triggered |
| 636 | # at offload block boundaries. When the prompt is exactly at a boundary, |
| 637 | # CPU offload must not load that boundary's cached state because it already |
| 638 | # includes the last token that must be recomputed. |
| 639 | # This is fine for other attention types as we have all the necessary |
| 640 | # token KV values in the hit blocks. |
| 641 | prompt = TokensPrompt(prompt_token_ids=initial_ids) |
| 642 | _verify(llm, prompt, "block-boundary-prompt") |
| 643 | |
| 644 | # Test for prompt token ids at non-block boundaries. |
| 645 | # Reuse is okay for this case. |
| 646 | prompt = TokensPrompt(prompt_token_ids=[0] + initial_ids) |
| 647 | _verify(llm, prompt, "block-mid-prompt") |
| 648 | |
| 649 | assert not failures, "\n\n".join(failures) |
| 650 | |
| 651 | finally: |
| 652 | del llm |