| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Tests for KV cache offloading configuration.""" |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.config import CacheConfig, KVTransferConfig, ParallelConfig, VllmConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 10 | |
| 11 | pytestmark = pytest.mark.cpu_test |
| 12 | |
| 13 | |
| 14 | class _StubLMCacheMPConnector: |
| 15 | """Stand-in for LMCacheMPConnector used in config-translation tests. |
| 16 | |
| 17 | The real connector module hard-imports the optional ``lmcache`` package |
| 18 | at module load time, which is not installed in the cpu_test image. This |
| 19 | test only asserts on the connector *name* and the ``extra_config`` dict |
| 20 | produced by ``VllmConfig``, never instantiates the connector, so a bare |
| 21 | placeholder class is sufficient. Not subclassing ``SupportsHMA`` mirrors |
| 22 | the real connector's HMA support (it does not support HMA either).""" |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def stub_lmcache_mp_connector(monkeypatch): |
| 27 | """Replace the lazy loader so VllmConfig.__post_init__ does not import |
| 28 | ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests.""" |
| 29 | monkeypatch.setitem( |
| 30 | KVConnectorFactory._registry, |
| 31 | "LMCacheMPConnector", |
| 32 | lambda: _StubLMCacheMPConnector, |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @pytest.mark.parametrize( |
| 37 | "kv_offloading_backend,kv_offloading_size,tp,pp,expected_backend,expected_bytes", |
| 38 | [ |
| 39 | ("native", 4.0, 1, 1, "OffloadingConnector", 4.0 * (1 << 30)), |
| 40 | # bytes per rank: 8.0 GiB / (2 * 2) = 2.0 GiB |
| 41 | ("native", 8.0, 2, 2, "OffloadingConnector", 8.0 * (1 << 30)), |
| 42 | # ``lmcache`` backend now defaults to LMCacheMPConnector. The KV |
| 43 | # storage capacity is owned by the standalone LMCache server, so |
| 44 | # ``kv_offloading_size`` is intentionally not propagated. |
| 45 | ("lmcache", 4.0, 1, 1, "LMCacheMPConnector", None), |
| 46 | ("lmcache", 8.0, 2, 2, "LMCacheMPConnector", None), |
| 47 | # When kv_offloading_size is None, offloading is disabled (backend is ignored) |
| 48 | ("native", None, 1, 1, None, None), |
| 49 | ], |
| 50 | ) |
| 51 | def test_kv_connector( |
| 52 | stub_lmcache_mp_connector, |
| 53 | kv_offloading_backend, |
| 54 | kv_offloading_size, |
| 55 | tp, |
| 56 | pp, |
| 57 | expected_backend, |
| 58 | expected_bytes, |
| 59 | ): |
| 60 | kv_transfer_config = ( |
| 61 | KVTransferConfig(kv_connector_extra_config={"existing_key": "existing_value"}) |
| 62 | if expected_backend is not None |
| 63 | else None |
| 64 | ) |
| 65 | |
| 66 | vllm_config = VllmConfig( |
| 67 | cache_config=CacheConfig( |
| 68 | kv_offloading_backend=kv_offloading_backend, |
| 69 | kv_offloading_size=kv_offloading_size, |
| 70 | ), |
| 71 | kv_transfer_config=kv_transfer_config, |
| 72 | parallel_config=ParallelConfig( |
| 73 | tensor_parallel_size=tp, pipeline_parallel_size=pp |
| 74 | ), |
| 75 | ) |
| 76 | |
| 77 | # No KV transfer config expected |
| 78 | if expected_backend is None: |
| 79 | assert vllm_config.kv_transfer_config is expected_backend |
| 80 | return |
| 81 | |
| 82 | kv_transfer_config = vllm_config.kv_transfer_config |
| 83 | kv_connector_extra_config = kv_transfer_config.kv_connector_extra_config |
| 84 | |
| 85 | assert kv_transfer_config.kv_connector == expected_backend |
| 86 | assert kv_transfer_config.kv_role == "kv_both" |
| 87 | |
| 88 | if kv_offloading_backend == "native": |
| 89 | assert kv_connector_extra_config["cpu_bytes_to_use"] == expected_bytes |
| 90 | # Existing config should be preserved |
| 91 | assert kv_connector_extra_config["existing_key"] == "existing_value" |
| 92 | elif kv_offloading_backend == "lmcache": |
| 93 | # MP mode does not push lmcache.local_cpu / max_local_cpu_size into |
| 94 | # extra config (the LMCache server owns capacity). Pre-existing |
| 95 | # extra config entries are preserved as-is. |
| 96 | assert "lmcache.local_cpu" not in kv_connector_extra_config |
| 97 | assert "lmcache.max_local_cpu_size" not in kv_connector_extra_config |
| 98 | assert kv_connector_extra_config["existing_key"] == "existing_value" |
| 99 | |
| 100 | |
| 101 | def _build_config( |
| 102 | *, |
| 103 | kv_connector: str | None, |
| 104 | enable_sleep_mode: bool = False, |
| 105 | enable_cumem_allocator: bool = False, |
| 106 | ) -> VllmConfig: |
| 107 | """Build a VllmConfig that exercises _verify_kv_transfer_compat without |
| 108 | requiring a real model (avoids HF downloads in CI).""" |
| 109 | from types import SimpleNamespace |
| 110 | |
| 111 | kv_transfer_config = ( |
| 112 | KVTransferConfig(kv_connector=kv_connector, kv_role="kv_both") |
| 113 | if kv_connector is not None |
| 114 | else None |
| 115 | ) |
| 116 | cfg = VllmConfig.__new__(VllmConfig) |
| 117 | cfg.kv_transfer_config = kv_transfer_config |
| 118 | cfg.model_config = SimpleNamespace( |
| 119 | enable_sleep_mode=enable_sleep_mode, |
| 120 | enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode), |
| 121 | ) |
| 122 | cfg._verify_kv_transfer_compat() |
| 123 | return cfg |
| 124 | |
| 125 | |
| 126 | @pytest.mark.parametrize( |
| 127 | "kv_connector", ["NixlConnector", "MooncakeConnectorV1", "SomeOOTConnector"] |
| 128 | ) |
| 129 | def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector): |
| 130 | """KV connectors that pin KV cache memory (e.g. via ibv_reg_mr) are |
| 131 | invalidated when expandable_segments lets the CUDA VMM allocator remap |
| 132 | the underlying physical pages. We can't enumerate every connector that |
| 133 | does this (especially OOT ones), so reject the combination whenever any |
| 134 | connector is configured.""" |
| 135 | monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| 136 | with pytest.raises(ValueError, match="expandable_segments"): |
| 137 | _build_config(kv_connector=kv_connector) |
| 138 | |
| 139 | |
| 140 | def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch): |
| 141 | """Sleep mode routes KV allocations through CuMemAllocator's pool, which |
| 142 | auto-disables expandable_segments (see #40812).""" |
| 143 | monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| 144 | _build_config(kv_connector="NixlConnector", enable_sleep_mode=True) |
| 145 | |
| 146 | |
| 147 | def test_kv_connector_allows_expandable_segments_with_cumem_allocator( |
| 148 | monkeypatch, |
| 149 | ): |
| 150 | """Manual CuMem allocation must also bypass expandable_segments.""" |
| 151 | monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| 152 | _build_config(kv_connector="NixlConnector", enable_cumem_allocator=True) |
| 153 | |
| 154 | |
| 155 | def test_kv_connector_allows_other_alloc_conf(monkeypatch): |
| 156 | """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected.""" |
| 157 | monkeypatch.setenv( |
| 158 | "PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:512,expandable_segments:False" |
| 159 | ) |
| 160 | _build_config(kv_connector="NixlConnector") |
| 161 | |
| 162 | |
| 163 | def test_no_kv_connector_ignores_expandable_segments(monkeypatch): |
| 164 | """The expandable_segments check only applies when a KV connector is |
| 165 | configured.""" |
| 166 | monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| 167 | _build_config(kv_connector=None) |
| 168 | |
| 169 | |
| 170 | def test_kv_offloading_size_only_uses_native_default(): |
| 171 | """Test that setting only kv_offloading_size enables native offloading.""" |
| 172 | vllm_config = VllmConfig( |
| 173 | cache_config=CacheConfig( |
| 174 | kv_offloading_size=4.0, |
| 175 | # kv_offloading_backend not set, should default to "native" |
| 176 | ), |
| 177 | ) |
| 178 | |
| 179 | kv_transfer_config = vllm_config.kv_transfer_config |
| 180 | kv_connector_extra_config = kv_transfer_config.kv_connector_extra_config |
| 181 | assert kv_transfer_config.kv_connector == "OffloadingConnector" |
| 182 | assert kv_transfer_config.kv_role == "kv_both" |
| 183 | assert kv_connector_extra_config["cpu_bytes_to_use"] == 4.0 * (1 << 30) |