iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Tier offload followup (#42529)

Browse Source 
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
This commit is contained in:
Ronen Schaffer
2026-05-18 19:41:58 +00:00
committed by GitHub
parent 8fc1c284b9
commit 84747489de
8 changed files with 226 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_offload/test_tiering_offloading.py
+23 -147
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for TieringOffloadingManager and ExampleSecondaryTier.
Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.
These tests verify:
1. Basic tiered offloading operations (store, load, lookup)
@@ -14,7 +14,6 @@ These tests verify:
from collections.abc import Iterable
from unittest.mock import MagicMock
import numpy as np
import pytest
import torch
@@ -23,8 +22,7 @@ from vllm.v1.kv_offload.base import (
ReqContext,
make_offload_key,
)
from vllm.v1.kv_offload.tiering.base import JobMetadata
from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
from vllm.v1.kv_offload.tiering.manager import (
CPUPrimaryTierOffloadingManager,
TieringOffloadingManager,
@@ -63,14 +61,17 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None:
return count
class TestExampleSecondaryTier:
"""Tests for ExampleSecondaryTier implementation."""
class TestExampleSecondaryTierManager:
"""Tests for ExampleSecondaryTierManager implementation."""
def test_basic_store_and_lookup(self):
"""Test basic store and lookup operations."""
mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
tier = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
tier = ExampleSecondaryTierManager(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
tier_type="example",
custom_param=67,
)
# Initially empty
@@ -88,81 +89,6 @@ class TestExampleSecondaryTier:
# Third block not present
assert tier.lookup(blocks[2], _CTX) is False
def test_lru_eviction(self):
"""Test LRU eviction policy."""
mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy())
tier = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=3
)
# Fill tier to capacity
blocks = to_keys(range(3))
for block in blocks:
tier.blocks[block] = True
assert tier.get_num_blocks() == 3
# Touch first block (make it most recently used)
tier.touch([blocks[0]], _CTX)
# Store new block should evict blocks[1] (least recently used)
new_block = to_keys([3])[0]
tier.submit_store(
JobMetadata(
job_id=1,
keys=[new_block],
block_ids=np.array([0], dtype=np.int64),
is_promotion=False,
req_context=_CTX,
)
)
# Complete the job
tier.get_finished()
# Verify new block is stored and blocks[1] was evicted (LRU)
assert new_block in tier.blocks
assert blocks[1] not in tier.blocks
# blocks[0] and blocks[2] should still be present
assert blocks[0] in tier.blocks
assert blocks[2] in tier.blocks
def test_async_simulation(self):
"""Test simulated async behavior."""
mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
tier = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
max_blocks=10,
simulate_async=True,
)
blocks = to_keys(range(2))
# Submit store job
tier.submit_store(
JobMetadata(
job_id=1,
keys=blocks,
block_ids=np.array([0, 1], dtype=np.int64),
is_promotion=False,
req_context=_CTX,
)
)
# Blocks should not yet be stored (pending async completion)
assert tier.get_num_blocks() == 0
# First get_finished() should complete the job
completed = list(tier.get_finished())
assert len(completed) == 1
assert completed[0].job_id == 1
assert completed[0].success is True
# Blocks should now be stored
assert tier.get_num_blocks() == 2
class TestTieringOffloadingManager:
"""Tests for TieringOffloadingManager."""
@@ -178,11 +104,15 @@ class TestTieringOffloadingManager:
mock_view = mock_region.create_kv_memoryview()
# Create secondary tiers with the primary view
self.secondary_tier1 = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
self.secondary_tier1 = ExampleSecondaryTierManager(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
tier_type="example",
)
self.secondary_tier2 = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
self.secondary_tier2 = ExampleSecondaryTierManager(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
tier_type="example",
)
# Create tiered manager
@@ -359,20 +289,19 @@ class TestTieringOffloadingManager:
self.manager.complete_store(blocks, _CTX, success=True)
list(self.manager.take_events())
self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
# Touch blocks
self.manager.touch(blocks, _CTX)
# Verify touch was called on primary tier (check LRU order)
# In LRU, touched blocks should be at the end
primary_keys = list(self.primary_tier._policy.blocks.keys())
assert primary_keys[-3:] == list(reversed(blocks))
# Verify touch was called on all secondary tiers
secondary1_keys = list(self.secondary_tier1.blocks.keys())
assert secondary1_keys[-3:] == list(reversed(blocks))
secondary2_keys = list(self.secondary_tier2.blocks.keys())
assert secondary2_keys[-3:] == list(reversed(blocks))
# Verify touch was propagated to all secondary tiers
self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX)
self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX)
def test_failed_store_no_cascade(self, manager_setup):
"""Test that failed GPU→primary store doesn't cascade."""
@@ -396,59 +325,6 @@ class TestTieringOffloadingManager:
self.secondary_tier1.submit_store.assert_not_called()
self.secondary_tier2.submit_store.assert_not_called()
def test_multiple_secondary_tiers_independent_eviction(self):
"""Test that secondary tiers manage their own evictions."""
mock_region = _mock_mmap_region(10)
mock_view = mock_region.create_kv_memoryview()
# Create tier with small capacity
small_tier = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
max_blocks=5,
simulate_async=False,
)
large_tier = ExampleSecondaryTier(
vllm_config=_MOCK_VLLM_CONFIG,
primary_kv_view=mock_view,
max_blocks=10,
simulate_async=False,
)
# Create a fresh primary tier for this test
primary_tier = CPUPrimaryTierOffloadingManager(
num_blocks=10, mmap_region=mock_region
)
manager = TieringOffloadingManager(
primary_tier=primary_tier,
secondary_tiers=[small_tier, large_tier],
)
# First, store 5 blocks to fill the small tier
blocks1 = to_keys(range(5))
result = manager.prepare_store(blocks1, _CTX)
assert result is not None
manager.complete_store(blocks1, _CTX, success=True)
list(manager.take_events())
# Both tiers should have 5 blocks
assert small_tier.get_num_blocks() == 5
assert large_tier.get_num_blocks() == 5
# Now store 3 more blocks - small tier should evict 3 blocks
blocks2 = to_keys(range(5, 8))
result = manager.prepare_store(blocks2, _CTX)
assert result is not None
manager.complete_store(blocks2, _CTX, success=True)
list(manager.take_events())
# Small tier should still have 5 blocks (evicted 3, added 3)
assert small_tier.get_num_blocks() == 5
# Large tier should have all 8 blocks
assert large_tier.get_num_blocks() == 8
def test_lookup_batches_submit_load_per_request(self, manager_setup):
"""lookup() defers submit_load until take_events(), one call per request.
vllm/v1/kv_offload/cpu/spec.py
+2 -4
View File
@@ -81,9 +81,7 @@ class CPUOffloadingSpec(OffloadingSpec):
)
return self._manager
def _create_handlers(
self, kv_caches: CanonicalKVCaches
) -> CpuGpuOffloadingHandlers:
def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
return CpuGpuOffloadingHandlers(
kv_caches=kv_caches,
block_size_factor=self.block_size_factor,
@@ -98,7 +96,7 @@ class CPUOffloadingSpec(OffloadingSpec):
raise Exception(
"CPU Offloading is currently only supported on CUDA-alike GPUs"
)
self._handlers = self._create_handlers(kv_caches)
self._handlers = self.create_handlers(kv_caches)
assert self._handlers is not None
yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler
vllm/v1/kv_offload/tiering/base.py
+14 -14
View File
@@ -53,9 +53,22 @@ class SecondaryTierManager(ABC):
async jobs; get_finished() polls for completion.
"""
def __init__(self, vllm_config: "VllmConfig", primary_kv_view: memoryview) -> None:
def __init__(
self,
vllm_config: "VllmConfig",
primary_kv_view: memoryview,
tier_type: str,
) -> None:
"""
Args:
vllm_config: Global vLLM configuration.
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
tier_type: Tier type identifier, set by SecondaryTierFactory
from the registered tier type.
"""
self._vllm_config = vllm_config
self._primary_kv_view: memoryview = primary_kv_view
self.tier_type = tier_type
@abstractmethod
def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
@@ -153,16 +166,3 @@ class SecondaryTierManager(ABC):
def shutdown(self) -> None:
"""Release resources held by this tier (threads, connections, etc.)."""
return
@staticmethod
@abstractmethod
def get_tier_type() -> str:
"""
Get the type identifier of this tier (e.g., "example", "storage").
Must match the "type" field in the tier config dict.
Returns:
Tier type string.
"""
pass
vllm/v1/kv_offload/tiering/example/__init__.py
-235
View File
@@ -1,235 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
ExampleSecondaryTier: A simple in-memory secondary tier for testing.
This implementation provides a minimal secondary tier that stores blocks
in memory (using a dictionary) and simulates async transfers with immediate
completion. It's useful for testing the TieringOffloadingManager without
requiring actual storage or network backends.
"""
from collections import OrderedDict
from collections.abc import Collection, Iterable
from dataclasses import dataclass
from typing import TYPE_CHECKING
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
from vllm.v1.kv_offload.tiering.base import (
JobId,
JobMetadata,
JobResult,
SecondaryTierManager,
)
if TYPE_CHECKING:
from vllm.config import VllmConfig
@dataclass
class _JobMetadata:
"""Internal metadata for tracking job details."""
job_id: JobId
keys: Collection[OffloadKey]
is_store: bool # True for store jobs, False for load jobs
class ExampleSecondaryTier(SecondaryTierManager):
"""
A simple in-memory secondary tier for testing.
This implementation:
- Stores blocks in a dictionary (key -> True)
- Simulates async transfers with immediate completion
- Uses LRU eviction policy
"""
def __init__(
self,
vllm_config: "VllmConfig",
primary_kv_view: memoryview,
max_blocks: int = 1000,
simulate_async: bool = False,
):
"""
Initialize the example secondary tier.
Args:
vllm_config: Global vLLM configuration.
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
max_blocks: Maximum number of blocks this tier can store
simulate_async: If True, jobs complete on next get_finished() call.
If False, jobs complete immediately.
"""
super().__init__(vllm_config, primary_kv_view)
self.max_blocks = max_blocks
self.simulate_async = simulate_async
# key -> True (only care about presence)
self.blocks: OrderedDict[OffloadKey, bool] = OrderedDict()
# Completed jobs waiting to be retrieved by get_finished()
self.completed_jobs: list[JobResult] = []
# Pending jobs (for simulated async mode)
self.pending_jobs: list[_JobMetadata] = []
def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
"""
Check whether a block exists in this secondary tier.
Args:
key: Offload key to look up.
req_context: Per-request context.
Returns:
True if the block is present, False if not found.
"""
return key in self.blocks
def submit_store(self, job_metadata: JobMetadata) -> None:
"""
Submit an async job to store blocks from primary tier to this tier.
Args:
job_metadata: Job metadata including job_id, keys, and
spec for reading blocks from the primary tier.
"""
job_id = job_metadata.job_id
keys = job_metadata.keys
block_ids = job_metadata.block_ids
assert len(keys) == len(block_ids), (
f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
)
# Filter out blocks already present
blocks_to_store = [bh for bh in keys if bh not in self.blocks]
if not blocks_to_store:
# All blocks already present
return
# Evict blocks if needed (LRU policy)
num_blocks_to_evict = len(blocks_to_store) - (
self.max_blocks - len(self.blocks)
)
evicted = []
if num_blocks_to_evict > 0:
# Collect eviction candidates first (LRU order), then delete atomically
protected = set(keys)
for key in self.blocks:
if key not in protected:
evicted.append(key)
if len(evicted) == num_blocks_to_evict:
break
else:
# Could not collect enough eviction candidates
return
for key in evicted:
del self.blocks[key]
# Create internal job metadata
internal_job_metadata = _JobMetadata(
job_id=job_id, keys=blocks_to_store, is_store=True
)
if self.simulate_async:
# Job will complete on next get_finished() call
self.pending_jobs.append(internal_job_metadata)
else:
# Job completes immediately
self._complete_store_job(internal_job_metadata)
def submit_load(self, job_metadata: JobMetadata) -> None:
"""
Submit an async job to load blocks from this tier to primary tier.
Args:
job_metadata: Job metadata including job_id, keys, and
spec for writing blocks into the primary tier.
"""
job_id = job_metadata.job_id
keys = job_metadata.keys
block_ids = job_metadata.block_ids
assert len(keys) == len(block_ids), (
f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
)
# Verify all blocks exist
for key in keys:
if key not in self.blocks:
return
# Create internal job metadata
internal_job_metadata = _JobMetadata(job_id=job_id, keys=keys, is_store=False)
if self.simulate_async:
# Job will complete on next get_finished() call
self.pending_jobs.append(internal_job_metadata)
else:
# Job completes immediately
self._complete_load_job(internal_job_metadata)
def get_finished(self) -> Iterable[JobResult]:
"""
Poll for finished async jobs.
Returns:
Iterable of JobResult objects for all jobs that have
finished since the last call.
"""
# Move pending jobs to completed
if self.simulate_async and self.pending_jobs:
for job_metadata in self.pending_jobs:
if job_metadata.is_store:
self._complete_store_job(job_metadata)
else:
self._complete_load_job(job_metadata)
self.pending_jobs.clear()
# Return completed jobs
result = self.completed_jobs
self.completed_jobs = []
return result
def _complete_store_job(self, job_metadata: _JobMetadata):
"""Complete a store job by adding blocks to storage."""
for key in job_metadata.keys:
self.blocks[key] = True
# Return simplified JobResult (only job_id and success)
self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
def _complete_load_job(self, job_metadata: _JobMetadata):
"""Complete a load job."""
# Return simplified JobResult (only job_id and success)
self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
"""
Mark blocks as recently used (move to end of LRU list).
Args:
keys: Blocks to mark as recently used.
req_context: Per-request context.
"""
for key in reversed(list(keys)):
if key in self.blocks:
self.blocks.move_to_end(key)
@staticmethod
def get_tier_type() -> str:
return "example"
def get_num_blocks(self) -> int:
"""Get the number of blocks currently stored in this tier."""
return len(self.blocks)
def clear(self):
"""Clear all blocks and pending jobs (for testing)."""
self.blocks.clear()
self.completed_jobs.clear()
self.pending_jobs.clear()
vllm/v1/kv_offload/tiering/example/manager.py
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
ExampleSecondaryTierManager: A simple in-memory secondary tier.
This implementation provides a minimal secondary tier that stores blocks
in memory (using a dictionary) with immediate completion. It serves as a
reference for writing new tiers and is useful for testing the
TieringOffloadingManager without requiring actual storage or network backends.
"""
import logging
from collections.abc import Iterable
from typing import TYPE_CHECKING
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
from vllm.v1.kv_offload.tiering.base import (
JobMetadata,
JobResult,
SecondaryTierManager,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from vllm.config import VllmConfig
class ExampleSecondaryTierManager(SecondaryTierManager):
"""
A simple in-memory secondary tier.
This implementation:
- Stores blocks in a dictionary (key -> True)
- Completes transfers immediately (synchronous)
"""
def __init__(
self,
vllm_config: "VllmConfig",
primary_kv_view: memoryview,
tier_type: str,
custom_param: int = 0,
):
"""
Initialize the example secondary tier.
Args:
custom_param: Dummy parameter demonstrating custom args.
"""
super().__init__(
vllm_config=vllm_config,
primary_kv_view=primary_kv_view,
tier_type=tier_type,
)
logger.info(
"ExampleSecondaryTierManager initialized with custom_param=%d", custom_param
)
# key -> True (only care about presence)
self.blocks: dict[OffloadKey, bool] = {}
# Completed jobs waiting to be retrieved by get_finished()
self.completed_jobs: list[JobResult] = []
def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
"""
Check whether a block exists in this secondary tier.
Args:
key: Offload key to look up.
req_context: Per-request context.
Returns:
True if the block is present, False if not found.
"""
return key in self.blocks
def submit_store(self, job_metadata: JobMetadata) -> None:
"""
Submit a job to store blocks from primary tier to this tier.
Args:
job_metadata: Job metadata including job_id, keys, and
spec for reading blocks from the primary tier.
"""
keys = job_metadata.keys
block_ids = job_metadata.block_ids
assert len(keys) == len(block_ids), (
f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
)
for key in keys:
self.blocks[key] = True
self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
def submit_load(self, job_metadata: JobMetadata) -> None:
"""
Submit a job to load blocks from this tier to primary tier.
Args:
job_metadata: Job metadata including job_id, keys, and
spec for writing blocks into the primary tier.
"""
keys = job_metadata.keys
block_ids = job_metadata.block_ids
assert len(keys) == len(block_ids), (
f"Length mismatch: {len(keys)} keys but {len(block_ids)} block_ids"
)
for key in keys:
if key not in self.blocks:
self.completed_jobs.append(
JobResult(job_id=job_metadata.job_id, success=False)
)
return
self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
def get_finished(self) -> Iterable[JobResult]:
"""
Poll for finished jobs.
Returns:
Iterable of JobResult objects for all jobs that have
finished since the last call.
"""
result = self.completed_jobs
self.completed_jobs = []
return result
def get_num_blocks(self) -> int:
"""Get the number of blocks currently stored in this tier."""
return len(self.blocks)
vllm/v1/kv_offload/tiering/factory.py
+41 -38
View File
@@ -1,56 +1,59 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Factory for creating secondary tier implementations.
"""
import importlib
from collections.abc import Callable
from typing import TYPE_CHECKING
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
if TYPE_CHECKING:
from vllm.config import VllmConfig
SUPPORTED_TIERS: tuple[type[SecondaryTierManager], ...] = (ExampleSecondaryTier,)
_TIER_REGISTRY: dict[str, type[SecondaryTierManager]] = {
cls.get_tier_type(): cls for cls in SUPPORTED_TIERS
}
class SecondaryTierFactory:
_registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {}
@classmethod
def register_tier(cls, tier_type: str, module_path: str, class_name: str) -> None:
if tier_type in cls._registry:
raise ValueError(f"Tier '{tier_type}' is already registered.")
def create_secondary_tier(
tier_config: dict,
primary_kv_view: memoryview,
vllm_config: "VllmConfig",
) -> SecondaryTierManager:
"""
Create a secondary tier from configuration.
def loader() -> type[SecondaryTierManager]:
module = importlib.import_module(module_path)
return getattr(module, class_name)
Args:
tier_config: Dictionary with tier configuration containing:
- type (required): Type of secondary tier (e.g., "example")
- Additional tier-specific parameters are passed directly
to the tier constructor
primary_kv_view: Memoryview of the primary tier's CPU KV cache.
vllm_config: Global vLLM configuration.
cls._registry[tier_type] = loader
Returns:
SecondaryTierManager instance
@classmethod
def create_secondary_tier(
cls,
tier_config: dict,
primary_kv_view: memoryview,
vllm_config: "VllmConfig",
) -> SecondaryTierManager:
config = tier_config.copy()
Raises:
ValueError: If tier type is unknown or configuration is invalid
"""
config = tier_config.copy()
tier_type = config.pop("type", None)
if not tier_type:
raise ValueError("Secondary tier configuration must include 'type'")
tier_type = config.pop("type", None)
if not tier_type:
raise ValueError("Secondary tier configuration must include 'type'")
if tier_type not in cls._registry:
raise ValueError(
f"Unknown secondary tier type: {tier_type!r}. "
f"Supported types: {list(cls._registry)}"
)
cls = _TIER_REGISTRY.get(tier_type)
if cls is None:
raise ValueError(
f"Unknown secondary tier type: {tier_type!r}. "
f"Supported types: {list(_TIER_REGISTRY)}"
tier_cls = cls._registry[tier_type]()
return tier_cls(
vllm_config=vllm_config,
primary_kv_view=primary_kv_view,
tier_type=tier_type,
**config,
)
return cls(vllm_config=vllm_config, primary_kv_view=primary_kv_view, **config)
SecondaryTierFactory.register_tier(
"example",
"vllm.v1.kv_offload.tiering.example.manager",
"ExampleSecondaryTierManager",
)
vllm/v1/kv_offload/tiering/manager.py
+1 -1
View File
@@ -194,7 +194,7 @@ class TieringOffloadingManager(OffloadingManager):
job_metadata = self._transfer_jobs.pop(job_id, None)
assert job_metadata is not None, (
f"Finished job_id {job_id} from tier #{i}"
f" ({tier.get_tier_type()}) not in _transfer_jobs"
f" ({tier.tier_type}) not in _transfer_jobs"
)
if job_metadata.is_promotion:
vllm/v1/kv_offload/tiering/spec.py
+8 -9
View File
@@ -25,15 +25,14 @@ Example configuration:
"secondary_tiers": [
{
"type": "example",
# Tier-specific parameters (for ExampleSecondaryTier):
"max_blocks": 10000,
"simulate_async": False
"custom_param": 67
}
]
}
"""
import torch
from typing_extensions import override
from vllm.config import VllmConfig
from vllm.logger import init_logger
@@ -42,7 +41,7 @@ from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager
from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
from vllm.v1.kv_offload.tiering.factory import create_secondary_tier
from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
from vllm.v1.kv_offload.tiering.manager import (
CPUPrimaryTierOffloadingManager,
TieringOffloadingManager,
@@ -77,6 +76,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
# Scheduler-side mmap (rank=None); kept for cleanup
self._scheduler_mmap: SharedOffloadRegion | None = None
@override
def get_manager(self) -> OffloadingManager:
"""
Get the TieringOffloadingManager.
@@ -123,14 +123,14 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
secondary_tiers = []
for i, tier_config in enumerate(self.secondary_tier_configs):
try:
tier = create_secondary_tier(
tier = SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, self.vllm_config
)
secondary_tiers.append(tier)
logger.info(
"Created secondary tier #%d (%s)",
i,
tier.get_tier_type(),
tier.tier_type,
)
except Exception as e:
logger.error(
@@ -164,9 +164,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
return self._manager
def _create_handlers(
self, kv_caches: CanonicalKVCaches
) -> CpuGpuOffloadingHandlers:
@override
def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
world_size = self.vllm_config.parallel_config.world_size
rank = torch.accelerator.current_device_index()
worker_mmap = SharedOffloadRegion(