Skip to content

vllm.v1.kv_offload.cpu.common

Classes:

CPULoadStoreSpec

Bases: BlockIDsLoadStoreSpec

Spec for loading/storing KV chunks to/from CPU memory.

The inherited block_ids field holds chunk indices into the CPU cache (not GPU block IDs). The chunk_ids alias exposes the same array under the name used by the tiering layer.

Source code in vllm/v1/kv_offload/cpu/common.py 
20
21
22
23
24
25
26
27
28
29
30
class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
    """Spec for loading/storing KV chunks to/from CPU memory.

    The inherited block_ids field holds chunk indices into the
    CPU cache (not GPU block IDs). The chunk_ids alias exposes
    the same array under the name used by the tiering layer.
    """

    @property
    def chunk_ids(self) -> np.ndarray:
        return self.block_ids

CPUOffloadingInfo dataclass

Bases: ConfigInfo

Static facts about the CPU offload tier.

Source code in vllm/v1/kv_offload/cpu/common.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@dataclass(frozen=True)
class CPUOffloadingInfo(ConfigInfo):
    """Static facts about the CPU offload tier."""

    # Chunk slots in the tier. Chunks, not GPU blocks.
    cpu_num_chunks: int
    # GPU blocks for each chunk.
    cpu_blocks_per_chunk: int
    # Size of the tier in bytes, including chunk padding.
    cpu_kv_cache_size_bytes: int
    # KV tokens the tier holds for requests of max_model_len tokens. None
    # when max_model_len is not known.
    cpu_kv_cache_size_tokens: int | None

    @classmethod
    def from_config(
        cls, config: OffloadingConfig, num_chunks: int, kv_bytes_per_chunk: int
    ) -> "CPUOffloadingInfo":
        return cls(
            cpu_num_chunks=num_chunks,
            cpu_blocks_per_chunk=config.cache.blocks_per_chunk,
            cpu_kv_cache_size_bytes=num_chunks * kv_bytes_per_chunk,
            cpu_kv_cache_size_tokens=_kv_cache_size_tokens(config, num_chunks),
        )

_kv_cache_size_tokens(config, num_chunks)

KV tokens the tier holds at max_model_len, like GPU kv_cache_size_tokens.

Source code in vllm/v1/kv_offload/cpu/common.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def _kv_cache_size_tokens(config: OffloadingConfig, num_chunks: int) -> int | None:
    """KV tokens the tier holds at max_model_len, like GPU kv_cache_size_tokens."""
    max_model_len = config.model.max_model_len
    if max_model_len <= 0 or not config.groups:
        return None
    chunks_per_request = 0
    for group in config.groups:
        group_chunks = cdiv(
            max_model_len, group.tokens_per_block * config.cache.blocks_per_chunk
        )
        if group.sliding_window_size_in_chunks is not None:
            group_chunks = min(group_chunks, group.sliding_window_size_in_chunks)
        chunks_per_request += group_chunks
    return num_chunks * max_model_len // chunks_per_request