A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 06:45 UTC
[KV Offload] Pass `OffloadingSpec` instead of `VllmConfig` to secondary tiers (#43076)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
May 20, 2026 at 06:32 UTC
4f940896a32c9e2a0eba7f50d521bf5f6b4de458
5 files changed
+15
-15
tests/v1/kv_offload/test_tiering_offloading.py
+4
-4
@@ -29,7 +29,7 @@ from vllm.v1.kv_offload.tiering.manager import (
29
)
30
31
_CTX = ReqContext(req_id="test")
32
-_MOCK_VLLM_CONFIG = MagicMock()
32
+_MOCK_OFFLOADING_SPEC = MagicMock()
33
34
35
def _mock_mmap_region(num_blocks: int, row_bytes: int = 16):
@@ -68,7 +68,7 @@ class TestExampleSecondaryTierManager:
68
"""Test basic store and lookup operations."""
69
mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
70
tier = ExampleSecondaryTierManager(
71
- vllm_config=_MOCK_VLLM_CONFIG,
71
+ offloading_spec=_MOCK_OFFLOADING_SPEC,
72
primary_kv_view=mock_view,
73
tier_type="example",
74
custom_param=67,
@@ -105,12 +105,12 @@ class TestTieringOffloadingManager:
105
106
# Create secondary tiers with the primary view
107
self.secondary_tier1 = ExampleSecondaryTierManager(
108
- vllm_config=_MOCK_VLLM_CONFIG,
108
+ offloading_spec=_MOCK_OFFLOADING_SPEC,
109
primary_kv_view=mock_view,
110
tier_type="example",
111
)
112
self.secondary_tier2 = ExampleSecondaryTierManager(
113
- vllm_config=_MOCK_VLLM_CONFIG,
113
+ offloading_spec=_MOCK_OFFLOADING_SPEC,
114
primary_kv_view=mock_view,
115
tier_type="example",
116
)
vllm/v1/kv_offload/tiering/base.py
+4
-4
@@ -14,7 +14,7 @@ import numpy as np
14
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
15
16
if TYPE_CHECKING:
17
- from vllm.config import VllmConfig
17
+ from vllm.v1.kv_offload.base import OffloadingSpec
18
19
# Type alias for job IDs used in async transfer tracking
20
JobId = int
@@ -55,18 +55,18 @@ class SecondaryTierManager(ABC):
55
56
def __init__(
57
self,
58
- vllm_config: "VllmConfig",
58
+ offloading_spec: "OffloadingSpec",
59
primary_kv_view: memoryview,
60
tier_type: str,
61
) -> None:
62
"""
63
Args:
64
- vllm_config: Global vLLM configuration.
64
+ offloading_spec: Offloading configuration.
65
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
66
tier_type: Tier type identifier, set by SecondaryTierFactory
67
from the registered tier type.
68
"""
69
- self._vllm_config = vllm_config
69
+ self._offloading_spec = offloading_spec
70
self._primary_kv_view: memoryview = primary_kv_view
71
self.tier_type = tier_type
72
vllm/v1/kv_offload/tiering/example/manager.py
+3
-3
@@ -23,7 +23,7 @@ from vllm.v1.kv_offload.tiering.base import (
23
logger = logging.getLogger(__name__)
24
25
if TYPE_CHECKING:
26
- from vllm.config import VllmConfig
26
+ from vllm.v1.kv_offload.base import OffloadingSpec
27
28
29
class ExampleSecondaryTierManager(SecondaryTierManager):
@@ -37,7 +37,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
37
38
def __init__(
39
self,
40
- vllm_config: "VllmConfig",
40
+ offloading_spec: "OffloadingSpec",
41
primary_kv_view: memoryview,
42
tier_type: str,
43
custom_param: int = 0,
@@ -49,7 +49,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
49
custom_param: Dummy parameter demonstrating custom args.
50
"""
51
super().__init__(
52
- vllm_config=vllm_config,
52
+ offloading_spec=offloading_spec,
53
primary_kv_view=primary_kv_view,
54
tier_type=tier_type,
55
)
vllm/v1/kv_offload/tiering/factory.py
+3
-3
@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
7
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
8
9
if TYPE_CHECKING:
10
- from vllm.config import VllmConfig
10
+ from vllm.v1.kv_offload.base import OffloadingSpec
11
12
13
class SecondaryTierFactory:
@@ -29,7 +29,7 @@ class SecondaryTierFactory:
29
cls,
30
tier_config: dict,
31
primary_kv_view: memoryview,
32
- vllm_config: "VllmConfig",
32
+ offloading_spec: "OffloadingSpec",
33
) -> SecondaryTierManager:
34
config = tier_config.copy()
35
@@ -45,7 +45,7 @@ class SecondaryTierFactory:
45
46
tier_cls = cls._registry[tier_type]()
47
return tier_cls(
48
- vllm_config=vllm_config,
48
+ offloading_spec=offloading_spec,
49
primary_kv_view=primary_kv_view,
50
tier_type=tier_type,
51
**config,
vllm/v1/kv_offload/tiering/spec.py
+1
-1
@@ -124,7 +124,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
124
for i, tier_config in enumerate(self.secondary_tier_configs):
125
try:
126
tier = SecondaryTierFactory.create_secondary_tier(
127
- tier_config, primary_kv_view, self.vllm_config
127
+ tier_config, primary_kv_view, self
128
)
129
secondary_tiers.append(tier)
130
logger.info(