| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Integration tests for SimpleCPUOffloadConnector with real models.""" |
| 4 | |
| 5 | import time |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm import LLM, SamplingParams, TokensPrompt |
| 10 | from vllm.config import KVTransferConfig |
| 11 | from vllm.platforms import current_platform |
| 12 | |
| 13 | if not current_platform.is_cuda_alike(): |
| 14 | pytest.skip("Requires CUDA or ROCm", allow_module_level=True) |
| 15 | |
| 16 | # Small models for default CI / local runs (accuracy only). |
| 17 | SMALL_MODELS = [ |
| 18 | "meta-llama/Llama-3.2-1B-Instruct", |
| 19 | "google/gemma-3-1b-it", |
| 20 | ] |
| 21 | |
| 22 | # Large models for optional perf runs only (slow to load and execute). |
| 23 | PERF_MODELS = [ |
| 24 | "meta-llama/Llama-3.1-8B", |
| 25 | "openai/gpt-oss-20b", |
| 26 | ] |
| 27 | |
| 28 | |
| 29 | def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM: |
| 30 | kv_transfer_config = KVTransferConfig( |
| 31 | kv_connector="SimpleCPUOffloadConnector", |
| 32 | kv_role="kv_both", |
| 33 | kv_connector_extra_config={ |
| 34 | "cpu_bytes_to_use": cpu_bytes_to_use, |
| 35 | "lazy_offload": lazy, |
| 36 | }, |
| 37 | ) |
| 38 | return LLM( |
| 39 | model=model, |
| 40 | kv_cache_memory_bytes=40 << 30, # 40 GiB |
| 41 | disable_hybrid_kv_cache_manager=False, |
| 42 | enable_prefix_caching=True, |
| 43 | kv_transfer_config=kv_transfer_config, |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | def _flush_gpu_cache(llm: LLM, sampling_params: SamplingParams, seed: int = 0): |
| 48 | """Generate enough filler requests to allocate the entire GPU KV cache. |
| 49 | |
| 50 | This pushes all prior blocks through the free queue so that the lazy |
| 51 | cursor offloads them to CPU before they are evicted. |
| 52 | """ |
| 53 | cache_config = llm.llm_engine.vllm_config.cache_config |
| 54 | num_gpu_blocks = cache_config.num_gpu_blocks |
| 55 | block_size = cache_config.block_size |
| 56 | # Use 1.2x GPU capacity to give the lazy cursor enough scheduling steps |
| 57 | # to walk past all target blocks near the tail of the free queue. |
| 58 | assert num_gpu_blocks is not None |
| 59 | total_tokens_needed = int(num_gpu_blocks * block_size * 1.5) |
| 60 | |
| 61 | # Use token-id prompts so each filler is unique (no prefix sharing). |
| 62 | # Split into multiple requests to stay under max_model_len. |
| 63 | max_tokens_per_req = 4096 |
| 64 | num_fillers = (total_tokens_needed + max_tokens_per_req - 1) // max_tokens_per_req |
| 65 | batch_size = 10 |
| 66 | for i in range(0, num_fillers, batch_size): |
| 67 | batch_end = min(i + batch_size, num_fillers) |
| 68 | filler_prompts = [] |
| 69 | for j in range(i, batch_end): |
| 70 | ids = [seed * num_fillers + j + 1] * max_tokens_per_req |
| 71 | filler_prompts.append(TokensPrompt(prompt_token_ids=ids)) |
| 72 | llm.generate(filler_prompts, sampling_params, use_tqdm=False) |
| 73 | |
| 74 | |
| 75 | def _accuracy_test(llm: LLM, lazy: bool = False): |
| 76 | """Verify that CPU-loaded KV produces correct output.""" |
| 77 | sampling_params = SamplingParams(max_tokens=1, temperature=0) |
| 78 | prompt = "hi " * 2000 + "Let's count to ten. One, two, three, " |
| 79 | |
| 80 | # Cold run — populate GPU cache and trigger CPU offload |
| 81 | cold_output = llm.generate(prompt, sampling_params, use_tqdm=False)[0] |
| 82 | |
| 83 | # CPU hit runs |
| 84 | test_count = 10 |
| 85 | success_count = 0 |
| 86 | expected = cold_output.outputs[0].text |
| 87 | for i in range(test_count): |
| 88 | if lazy: |
| 89 | _flush_gpu_cache(llm, sampling_params, seed=i) |
| 90 | time.sleep(2) # let engine core drain pending transfers |
| 91 | |
| 92 | # Reset GPU prefix cache so next run must load from CPU |
| 93 | if not llm.reset_prefix_cache(): |
| 94 | print(f"GPU prefix cache reset failed for iteration {i}") |
| 95 | |
| 96 | output = llm.generate(prompt, sampling_params, use_tqdm=False)[0] |
| 97 | if output.outputs[0].text == expected: |
| 98 | success_count += 1 |
| 99 | |
| 100 | assert success_count >= 0.5 * test_count, ( |
| 101 | f"Accuracy too low: {success_count}/{test_count} matched '{expected}'" |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | def _latency_test(llm: LLM, lazy: bool = False): |
| 106 | """Verify CPU cache hit is faster than cold compute.""" |
| 107 | sampling_params = SamplingParams(max_tokens=1, seed=42) |
| 108 | prompt_token_ids = [0] * 10001 |
| 109 | |
| 110 | num_times_cpu_better = 0 |
| 111 | num_tests = 10 |
| 112 | for i in range(num_tests): |
| 113 | prompt_token_ids[0] = i |
| 114 | prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)] |
| 115 | |
| 116 | # Cold |
| 117 | time.sleep(2) # let engine core drain pending transfers |
| 118 | if not llm.reset_prefix_cache(): |
| 119 | print(f"GPU prefix cache reset failed for iteration {i}") |
| 120 | start = time.time() |
| 121 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 122 | cold_time = time.time() - start |
| 123 | |
| 124 | if lazy: |
| 125 | _flush_gpu_cache(llm, sampling_params, seed=i) |
| 126 | else: |
| 127 | # Eager mode: GPU hit ensures store completion is processed. |
| 128 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 129 | |
| 130 | time.sleep(2) # let engine core drain pending transfers |
| 131 | if not llm.reset_prefix_cache(): |
| 132 | print(f"GPU prefix cache reset failed for iteration {i}") |
| 133 | |
| 134 | # CPU hit |
| 135 | start = time.time() |
| 136 | llm.generate(prompts, sampling_params, use_tqdm=False) |
| 137 | cpu_time = time.time() - start |
| 138 | |
| 139 | if cpu_time < cold_time: |
| 140 | num_times_cpu_better += 1 |
| 141 | |
| 142 | assert num_times_cpu_better >= 0.8 * num_tests, ( |
| 143 | f"CPU hit only faster {num_times_cpu_better}/{num_tests} times" |
| 144 | ) |
| 145 | |
| 146 | |
| 147 | @pytest.mark.optional |
| 148 | @pytest.mark.slow_test |
| 149 | @pytest.mark.parametrize("model", SMALL_MODELS) |
| 150 | def test_simple_cpu_offload_accuracy(model: str): |
| 151 | """Store to CPU, reset GPU, load from CPU; verify output matches baseline.""" |
| 152 | llm = _make_llm(model, False, 1 << 30) # 1GB |
| 153 | try: |
| 154 | _accuracy_test(llm, lazy=False) |
| 155 | finally: |
| 156 | del llm |
| 157 | |
| 158 | |
| 159 | @pytest.mark.optional |
| 160 | @pytest.mark.slow_test |
| 161 | @pytest.mark.parametrize("model", PERF_MODELS) |
| 162 | def test_simple_cpu_offload_perf_latency(model: str): |
| 163 | """CPU KV hit should beat cold prefill on long context (large models only).""" |
| 164 | llm = _make_llm(model, False, 10 << 30) # 10GB |
| 165 | try: |
| 166 | _latency_test(llm, lazy=False) |
| 167 | finally: |
| 168 | del llm |
| 169 | |
| 170 | |
| 171 | @pytest.mark.optional |
| 172 | @pytest.mark.slow_test |
| 173 | @pytest.mark.parametrize("model", SMALL_MODELS) |
| 174 | def test_simple_cpu_offload_accuracy_lazy(model: str): |
| 175 | """Lazy mode: flush GPU cache to trigger CPU offload, then verify hit.""" |
| 176 | # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks. |
| 177 | llm = _make_llm(model, True, 80 << 30) # 80GB |
| 178 | try: |
| 179 | _accuracy_test(llm, lazy=True) |
| 180 | finally: |
| 181 | del llm |
| 182 | |
| 183 | |
| 184 | @pytest.mark.optional |
| 185 | @pytest.mark.slow_test |
| 186 | @pytest.mark.parametrize("model", PERF_MODELS) |
| 187 | def test_simple_cpu_offload_perf_latency_lazy(model: str): |
| 188 | """Lazy mode: CPU KV hit should beat cold prefill (large models only).""" |
| 189 | # CPU must be larger than GPU KV cache to avoid evicting offloaded blocks. |
| 190 | llm = _make_llm(model, True, 80 << 30) # 80GB |
| 191 | try: |
| 192 | _latency_test(llm, lazy=True) |
| 193 | finally: |
| 194 | del llm |