Skip to content

vllm.v1.kv_offload.cpu

CPUOffloadingSpec

Bases: OffloadingSpec

Source code in vllm/v1/kv_offload/cpu.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class CPUOffloadingSpec(OffloadingSpec):

    def __init__(self, vllm_config: VllmConfig):
        super().__init__(vllm_config)

        num_cpu_blocks = self.extra_config.get("num_cpu_blocks")
        if not num_cpu_blocks:
            raise Exception("num_cpu_blocks must be specified "
                            "in kv_connector_extra_config")
        self.num_cpu_blocks: int = num_cpu_blocks

        # scheduler-side
        self._manager: Optional[OffloadingManager] = None

        # worker-side
        self._handler: Optional[OffloadingHandler] = None

    def get_manager(self) -> OffloadingManager:
        if not self._manager:
            kv_events_config = self.vllm_config.kv_events_config
            enable_events = (kv_events_config is not None
                             and kv_events_config.enable_kv_cache_events)
            self._manager = LRUOffloadingManager(CPUBackend(
                block_size=self.offloaded_block_size,
                num_blocks=self.num_cpu_blocks),
                                                 enable_events=enable_events)
        return self._manager

    def get_handlers(
        self, kv_caches: dict[str, torch.Tensor]
    ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec],
                        OffloadingHandler]]:
        if not self._handler:
            if not current_platform.is_cuda():
                raise Exception("CPU Offloading is currently only supported"
                                " on CUDA GPUs")

            layer_names = list(kv_caches.keys())
            layers = get_layers_from_vllm_config(self.vllm_config,
                                                 AttentionLayerBase,
                                                 layer_names)
            attn_backends = {
                layer_name: layers[layer_name].get_attn_backend()
                for layer_name in layer_names
            }

            self._handler = CpuGpuOffloadingHandler(
                attn_backends=attn_backends,
                gpu_block_size=self.gpu_block_size,
                cpu_block_size=self.offloaded_block_size,
                num_cpu_blocks=self.num_cpu_blocks,
                gpu_caches=kv_caches)

        assert self._handler is not None
        yield GPULoadStoreSpec, CPULoadStoreSpec, self._handler
        yield CPULoadStoreSpec, GPULoadStoreSpec, self._handler

_handler instance-attribute 

_handler: Optional[OffloadingHandler] = None

_manager instance-attribute 

_manager: Optional[OffloadingManager] = None

num_cpu_blocks instance-attribute 

num_cpu_blocks: int = num_cpu_blocks

__init__ 

__init__(vllm_config: VllmConfig)
Source code in vllm/v1/kv_offload/cpu.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def __init__(self, vllm_config: VllmConfig):
    super().__init__(vllm_config)

    num_cpu_blocks = self.extra_config.get("num_cpu_blocks")
    if not num_cpu_blocks:
        raise Exception("num_cpu_blocks must be specified "
                        "in kv_connector_extra_config")
    self.num_cpu_blocks: int = num_cpu_blocks

    # scheduler-side
    self._manager: Optional[OffloadingManager] = None

    # worker-side
    self._handler: Optional[OffloadingHandler] = None

get_handlers 

get_handlers(
    kv_caches: dict[str, Tensor],
) -> Iterator[
    tuple[
        type[LoadStoreSpec],
        type[LoadStoreSpec],
        OffloadingHandler,
    ]
]
Source code in vllm/v1/kv_offload/cpu.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def get_handlers(
    self, kv_caches: dict[str, torch.Tensor]
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec],
                    OffloadingHandler]]:
    if not self._handler:
        if not current_platform.is_cuda():
            raise Exception("CPU Offloading is currently only supported"
                            " on CUDA GPUs")

        layer_names = list(kv_caches.keys())
        layers = get_layers_from_vllm_config(self.vllm_config,
                                             AttentionLayerBase,
                                             layer_names)
        attn_backends = {
            layer_name: layers[layer_name].get_attn_backend()
            for layer_name in layer_names
        }

        self._handler = CpuGpuOffloadingHandler(
            attn_backends=attn_backends,
            gpu_block_size=self.gpu_block_size,
            cpu_block_size=self.offloaded_block_size,
            num_cpu_blocks=self.num_cpu_blocks,
            gpu_caches=kv_caches)

    assert self._handler is not None
    yield GPULoadStoreSpec, CPULoadStoreSpec, self._handler
    yield CPULoadStoreSpec, GPULoadStoreSpec, self._handler

get_manager 

get_manager() -> OffloadingManager
Source code in vllm/v1/kv_offload/cpu.py 
37
38
39
40
41
42
43
44
45
46
def get_manager(self) -> OffloadingManager:
    if not self._manager:
        kv_events_config = self.vllm_config.kv_events_config
        enable_events = (kv_events_config is not None
                         and kv_events_config.enable_kv_cache_events)
        self._manager = LRUOffloadingManager(CPUBackend(
            block_size=self.offloaded_block_size,
            num_blocks=self.num_cpu_blocks),
                                             enable_events=enable_events)
    return self._manager