[Test][KV Offloading] Add unit tests for OffloadingSpecFactory and SecondaryTierFactory (#46355)
Signed-off-by: Alex <alex.tech.lab@outlook.com>
AlexHuang committed
Jun 22, 2026 at 22:45 UTC
9a938df64e5e296e42524fc0a67956923806197d
4 files changed
+424
tests/v1/kv_offload/__init__.py
tests/v1/kv_offload/test_factory.py
new
+272
@@ -0,0 +1,272 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for OffloadingSpecFactory.
5
+
6
+These tests verify:
7
+1. Pre-registration integrity — registered module paths can actually import
8
+ and yield correct OffloadingSpec subclasses (CI sentinel against file moves).
9
+2. End-to-end factory → spec construction with real configs.
10
+3. Downstream collaboration — build_metric_definitions delegation.
11
+4. Error paths — unregistered specs, missing config, duplicate registration.
12
+"""
13
+
14
+import pytest
15
+import torch
16
+
17
+from vllm.config import KVTransferConfig
18
+from vllm.v1.kv_cache_interface import (
19
+ FullAttentionSpec,
20
+ KVCacheConfig,
21
+ KVCacheGroupSpec,
22
+ KVCacheTensor,
23
+)
24
+from vllm.v1.kv_offload.base import OffloadingSpec
25
+from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
26
+from vllm.v1.kv_offload.factory import OffloadingSpecFactory
27
+from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
28
+
29
+# ---------------------------------------------------------------------------
30
+# Fixtures
31
+# ---------------------------------------------------------------------------
32
+
33
+
34
+@pytest.fixture(autouse=True)
35
+def restore_registry():
36
+ """Save and restore OffloadingSpecFactory._registry between tests."""
37
+ original = dict(OffloadingSpecFactory._registry)
38
+ yield
39
+ OffloadingSpecFactory._registry = original
40
+
41
+
42
+def _make_vllm_config(
43
+ spec_name: str | None = "CPUOffloadingSpec",
44
+ cpu_bytes_to_use: int | None = None,
45
+ store_threshold: int = 0,
46
+ extra_config: dict | None = None,
47
+):
48
+ """Build a real VllmConfig with kv_transfer_config set for offloading."""
49
+ from vllm.config import (
50
+ CacheConfig,
51
+ DeviceConfig,
52
+ ModelConfig,
53
+ SchedulerConfig,
54
+ VllmConfig,
55
+ )
56
+
57
+ model_config = ModelConfig(
58
+ model="facebook/opt-125m",
59
+ trust_remote_code=True,
60
+ dtype="float16",
61
+ seed=42,
62
+ )
63
+ scheduler_config = SchedulerConfig(
64
+ max_num_seqs=16,
65
+ max_num_batched_tokens=64,
66
+ max_model_len=10000,
67
+ enable_chunked_prefill=True,
68
+ is_encoder_decoder=model_config.is_encoder_decoder,
69
+ )
70
+ cache_config = CacheConfig(
71
+ block_size=16,
72
+ gpu_memory_utilization=0.9,
73
+ cache_dtype="auto",
74
+ enable_prefix_caching=True,
75
+ )
76
+
77
+ cfg = extra_config or {}
78
+ if cpu_bytes_to_use is not None:
79
+ cfg["cpu_bytes_to_use"] = cpu_bytes_to_use
80
+ cfg["spec_name"] = spec_name
81
+ if store_threshold > 0:
82
+ cfg["store_threshold"] = store_threshold
83
+
84
+ kv_transfer_config = KVTransferConfig(
85
+ kv_connector="OffloadingConnector",
86
+ kv_role="kv_both",
87
+ kv_connector_extra_config=cfg,
88
+ )
89
+ return VllmConfig(
90
+ scheduler_config=scheduler_config,
91
+ model_config=model_config,
92
+ cache_config=cache_config,
93
+ kv_transfer_config=kv_transfer_config,
94
+ device_config=DeviceConfig("cpu"),
95
+ )
96
+
97
+
98
+def _make_kv_cache_config():
99
+ """Build a minimal KVCacheConfig with one KV cache tensor."""
100
+ num_blocks = 16
101
+ num_kv_heads = 1
102
+ head_size = 1
103
+ dtype = torch.float32
104
+ page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
105
+ kv_tensor = KVCacheTensor(
106
+ size=num_blocks * page_size, shared_by=["layer"], block_stride=0
107
+ )
108
+ return KVCacheConfig(
109
+ num_blocks=num_blocks,
110
+ kv_cache_tensors=[kv_tensor],
111
+ kv_cache_groups=[
112
+ KVCacheGroupSpec(
113
+ ["layer"],
114
+ FullAttentionSpec(
115
+ block_size=16,
116
+ num_kv_heads=num_kv_heads,
117
+ head_size=head_size,
118
+ dtype=dtype,
119
+ ),
120
+ )
121
+ ],
122
+ )
123
+
124
+
125
+# ---------------------------------------------------------------------------
126
+# Pre-registration integrity (CI sentinel)
127
+# ---------------------------------------------------------------------------
128
+
129
+
130
+def test_pre_registered_specs_can_be_imported():
131
+ """If someone moves cpu/spec.py but forgets to update factory.py, CI fails."""
132
+ for name in OffloadingSpecFactory._registry:
133
+ cls = OffloadingSpecFactory._registry[name]()
134
+ assert issubclass(cls, OffloadingSpec)
135
+
136
+
137
+def test_cpu_spec_registered():
138
+ """CPUOffloadingSpec is registered and importable."""
139
+ cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]()
140
+ assert cls is CPUOffloadingSpec
141
+
142
+
143
+def test_tiering_spec_registered():
144
+ """TieringOffloadingSpec is registered and importable."""
145
+ cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]()
146
+ assert cls is TieringOffloadingSpec
147
+
148
+
149
+# ---------------------------------------------------------------------------
150
+# Normal path — get_spec_cls
151
+# ---------------------------------------------------------------------------
152
+
153
+
154
+def test_get_spec_cls_returns_registered_class():
155
+ """Registered spec_name returns correct class."""
156
+ config = _make_vllm_config(spec_name="CPUOffloadingSpec")
157
+ spec_cls = OffloadingSpecFactory.get_spec_cls(config)
158
+ assert spec_cls is CPUOffloadingSpec
159
+
160
+
161
+def test_get_spec_cls_default_to_cpu():
162
+ """Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
163
+ config = _make_vllm_config(spec_name=None)
164
+ config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
165
+ spec_cls = OffloadingSpecFactory.get_spec_cls(config)
166
+ assert spec_cls is CPUOffloadingSpec
167
+
168
+
169
+# ---------------------------------------------------------------------------
170
+# End-to-end — create_spec
171
+# ---------------------------------------------------------------------------
172
+
173
+
174
+def test_create_cpu_offloading_spec_end_to_end():
175
+ """Full factory → spec construction with real VllmConfig/KVCacheConfig.
176
+
177
+ Verifies:
178
+ - cpu_bytes_to_use validation and num_blocks calculation
179
+ - block_size % hash_block_size assertion
180
+ - spec instance is CPUOffloadingSpec
181
+ """
182
+ config = _make_vllm_config(cpu_bytes_to_use=65536)
183
+ kv_cache_config = _make_kv_cache_config()
184
+ spec = OffloadingSpecFactory.create_spec(config, kv_cache_config)
185
+ assert isinstance(spec, CPUOffloadingSpec)
186
+ assert spec.num_blocks > 0
187
+
188
+
189
+# ---------------------------------------------------------------------------
190
+# Dynamic import via spec_module_path
191
+# ---------------------------------------------------------------------------
192
+
193
+
194
+def test_dynamic_load_via_spec_module_path():
195
+ """External spec loaded via spec_module_path.
196
+
197
+ This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec)
198
+ integrate with vLLM without being pre-registered in the factory.
199
+ The fallback path: registry miss → spec_module_path → importlib.import_module.
200
+ """
201
+ config = _make_vllm_config(spec_name="CPUOffloadingSpec")
202
+ # Delete from registry to force the dynamic import path
203
+ del OffloadingSpecFactory._registry["CPUOffloadingSpec"]
204
+ # spec_name not in registry → falls through to spec_module_path
205
+ config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
206
+ "vllm.v1.kv_offload.cpu.spec"
207
+ )
208
+ spec_cls = OffloadingSpecFactory.get_spec_cls(config)
209
+ assert spec_cls is CPUOffloadingSpec
210
+
211
+
212
+# ---------------------------------------------------------------------------
213
+# Error paths
214
+# ---------------------------------------------------------------------------
215
+
216
+
217
+def test_unregistered_spec_without_module_path_raises():
218
+ """spec_name not in registry + no spec_module_path → ValueError."""
219
+ config = _make_vllm_config(spec_name="NonexistentSpec")
220
+ with pytest.raises(ValueError, match="Unsupported spec type"):
221
+ OffloadingSpecFactory.get_spec_cls(config)
222
+
223
+ # create_spec should also fail (calls get_spec_cls internally)
224
+ kv_cache_config = _make_kv_cache_config()
225
+ with pytest.raises(ValueError, match="Unsupported spec type"):
226
+ OffloadingSpecFactory.create_spec(config, kv_cache_config)
227
+
228
+
229
+def test_cpu_spec_missing_cpu_bytes_to_use_raises():
230
+ """CPUOffloadingSpec requires cpu_bytes_to_use → Exception."""
231
+ config = _make_vllm_config(cpu_bytes_to_use=None)
232
+ config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
233
+ kv_cache_config = _make_kv_cache_config()
234
+ with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
235
+ OffloadingSpecFactory.create_spec(config, kv_cache_config)
236
+
237
+
238
+def test_duplicate_registration_raises():
239
+ """register_spec with existing name → ValueError."""
240
+ with pytest.raises(ValueError, match="is already registered"):
241
+ OffloadingSpecFactory.register_spec(
242
+ "CPUOffloadingSpec", "some.module", "SomeClass"
243
+ )
244
+
245
+
246
+# ---------------------------------------------------------------------------
247
+# Downstream collaboration — build_metric_definitions
248
+# ---------------------------------------------------------------------------
249
+
250
+
251
+def test_build_metric_definitions_empty_below_threshold():
252
+ """store_threshold < 2 → only base metric (no stores_skipped)."""
253
+ from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
254
+
255
+ config = _make_vllm_config(store_threshold=1)
256
+ spec_cls = OffloadingSpecFactory.get_spec_cls(config)
257
+ metrics = spec_cls.build_metric_definitions(
258
+ config.kv_transfer_config.kv_connector_extra_config
259
+ )
260
+ assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
261
+
262
+
263
+def test_build_metric_definitions_returns_counter_at_threshold():
264
+ """store_threshold >= 2 → returns stores_skipped counter definition."""
265
+ from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
266
+
267
+ config = _make_vllm_config(store_threshold=2)
268
+ spec_cls = OffloadingSpecFactory.get_spec_cls(config)
269
+ metrics = spec_cls.build_metric_definitions(
270
+ config.kv_transfer_config.kv_connector_extra_config
271
+ )
272
+ assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
tests/v1/kv_offload/tiering/__init__.py
tests/v1/kv_offload/tiering/test_factory.py
new
+152
@@ -0,0 +1,152 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for SecondaryTierFactory.
5
+
6
+These tests verify:
7
+1. Pre-registration integrity — registered tier module paths can import
8
+ and yield correct SecondaryTierManager subclasses (CI sentinel).
9
+2. Multi-tier creation via factory with correct tier_type propagation.
10
+3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
11
+"""
12
+
13
+from unittest.mock import MagicMock
14
+
15
+import pytest
16
+
17
+from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
18
+from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
19
+from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
20
+
21
+# ---------------------------------------------------------------------------
22
+# Fixtures
23
+# ---------------------------------------------------------------------------
24
+
25
+
26
+@pytest.fixture(autouse=True)
27
+def restore_registry():
28
+ """Save and restore SecondaryTierFactory._registry between tests."""
29
+ original = dict(SecondaryTierFactory._registry)
30
+ yield
31
+ SecondaryTierFactory._registry = original
32
+
33
+
34
+def _make_mock_args():
35
+ """Build common mock args for create_secondary_tier."""
36
+ return MagicMock(), MagicMock() # primary_kv_view, offloading_spec
37
+
38
+
39
+# ---------------------------------------------------------------------------
40
+# Pre-registration integrity (CI sentinel)
41
+# ---------------------------------------------------------------------------
42
+
43
+
44
+def test_pre_registered_tiers_can_be_imported():
45
+ """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager."""
46
+ for tier_type in SecondaryTierFactory._registry:
47
+ cls = SecondaryTierFactory._registry[tier_type]()
48
+ assert issubclass(cls, SecondaryTierManager)
49
+
50
+
51
+def test_example_tier_registered():
52
+ """Example tier is registered."""
53
+ cls = SecondaryTierFactory._registry["example"]()
54
+ assert cls is ExampleSecondaryTierManager
55
+
56
+
57
+# ---------------------------------------------------------------------------
58
+# Normal path — create_secondary_tier
59
+# ---------------------------------------------------------------------------
60
+
61
+
62
+def test_create_tier_from_registry():
63
+ """Registered tier_type creates instance with correct tier_type."""
64
+ primary_kv_view, offloading_spec = _make_mock_args()
65
+ tier_config = {"type": "example"}
66
+
67
+ tier = SecondaryTierFactory.create_secondary_tier(
68
+ tier_config, primary_kv_view, offloading_spec
69
+ )
70
+
71
+ assert isinstance(tier, SecondaryTierManager)
72
+ assert tier.tier_type == "example"
73
+
74
+
75
+def test_create_multiple_tiers():
76
+ """Multiple tier configs can be created with correct tier_types."""
77
+ primary_kv_view, offloading_spec = _make_mock_args()
78
+ configs = [
79
+ {"type": "example", "custom_param": 1},
80
+ {"type": "example", "custom_param": 2},
81
+ ]
82
+
83
+ tiers = [
84
+ SecondaryTierFactory.create_secondary_tier(
85
+ cfg.copy(), primary_kv_view, offloading_spec
86
+ )
87
+ for cfg in configs
88
+ ]
89
+
90
+ assert len(tiers) == 2
91
+ assert all(tier.tier_type == "example" for tier in tiers)
92
+ assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers)
93
+
94
+
95
+def test_register_new_tier_type():
96
+ """Verify that new tier types can be registered and created.
97
+
98
+ This is how external projects add custom secondary tiers
99
+ (e.g., llm-d FS backend was upstreamed as "fs" tier via this mechanism).
100
+ """
101
+ # Register a new tier type (reuse example manager for simplicity)
102
+ SecondaryTierFactory.register_tier(
103
+ "custom_tier",
104
+ "vllm.v1.kv_offload.tiering.example.manager",
105
+ "ExampleSecondaryTierManager",
106
+ )
107
+
108
+ primary_kv_view, offloading_spec = _make_mock_args()
109
+ tier = SecondaryTierFactory.create_secondary_tier(
110
+ {"type": "custom_tier", "custom_param": 99},
111
+ primary_kv_view,
112
+ offloading_spec,
113
+ )
114
+
115
+ assert tier.tier_type == "custom_tier"
116
+ assert isinstance(tier, ExampleSecondaryTierManager)
117
+
118
+
119
+# ---------------------------------------------------------------------------
120
+# Error paths
121
+# ---------------------------------------------------------------------------
122
+
123
+
124
+def test_missing_tier_type_raises():
125
+ """tier_config without 'type' → ValueError."""
126
+ primary_kv_view, offloading_spec = _make_mock_args()
127
+ tier_config: dict[str, str] = {}
128
+
129
+ with pytest.raises(ValueError, match="must include 'type'"):
130
+ SecondaryTierFactory.create_secondary_tier(
131
+ tier_config, primary_kv_view, offloading_spec
132
+ )
133
+
134
+
135
+def test_unknown_tier_type_raises():
136
+ """Unrecognized tier_type → ValueError with supported types list."""
137
+ primary_kv_view, offloading_spec = _make_mock_args()
138
+ tier_config = {"type": "nonexistent_tier"}
139
+
140
+ with pytest.raises(
141
+ ValueError,
142
+ match=r"Unknown secondary tier type.*Supported types:",
143
+ ):
144
+ SecondaryTierFactory.create_secondary_tier(
145
+ tier_config, primary_kv_view, offloading_spec
146
+ )
147
+
148
+
149
+def test_duplicate_registration_raises():
150
+ """register_tier with existing type → ValueError."""
151
+ with pytest.raises(ValueError, match="is already registered"):
152
+ SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")