mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 01:45:01 +08:00
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from vllm import LLM, SamplingParams
|
|
from vllm.config import KVTransferConfig
|
|
|
|
CPU_BLOCK_SIZES = [16, 48]
|
|
|
|
|
|
@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
|
|
def test_cpu_offloading(cpu_block_size: int) -> None:
|
|
"""
|
|
Tests OffloadingConnector with CPUOffloadingSpec.
|
|
"""
|
|
|
|
# configure OffloadingConnector (spec_name=CPUOffloadingSpec by default)
|
|
kv_transfer_config = KVTransferConfig(
|
|
kv_connector="OffloadingConnector",
|
|
kv_role="kv_both",
|
|
kv_connector_extra_config={
|
|
"num_cpu_blocks": 100,
|
|
"block_size": cpu_block_size
|
|
},
|
|
)
|
|
|
|
llm = LLM(
|
|
model="meta-llama/Llama-3.2-1B-Instruct",
|
|
gpu_memory_utilization=0.5,
|
|
kv_transfer_config=kv_transfer_config,
|
|
)
|
|
|
|
prompts = ["Hi " * 100]
|
|
sampling_params = SamplingParams(temperature=0, max_tokens=20)
|
|
|
|
# run generation - this should trigger saving KV cache
|
|
start_time = time.time()
|
|
llm.generate(prompts, sampling_params, use_tqdm=False)
|
|
cold_time = time.time() - start_time
|
|
|
|
# run generation again - should hit the GPU prefix cache
|
|
start_time = time.time()
|
|
llm.generate(prompts, sampling_params, use_tqdm=False)
|
|
gpu_hit_time = time.time() - start_time
|
|
|
|
# reset prefix cache to avoid GPU hit.
|
|
llm.reset_prefix_cache()
|
|
|
|
# sleep for a sec to make sure CPU finished storing
|
|
time.sleep(1)
|
|
|
|
# run generation again - this should trigger loading from CPU
|
|
start_time = time.time()
|
|
llm.generate(prompts, sampling_params, use_tqdm=False)
|
|
cpu_hit_time = time.time() - start_time
|
|
|
|
print("Generation times:")
|
|
print(f" Cold: {cold_time * 1000:.2f}ms")
|
|
print(f" GPU hit: {gpu_hit_time * 1000:.2f}ms")
|
|
print(f" CPU hit: {cpu_hit_time * 1000:.2f}ms")
|