Skip to content

vllm.v1.kv_offload.cpu.host_register

Registration of existing host memory for pinned DMA transfers.

Host allocations that the offloading workers share across processes (the SharedOffloadRegion mmap) cannot be allocated through the accelerator's pinned allocator, so they must be registered after the fact. Pinning is a pure bandwidth optimization: when registration is unavailable the transfers still run correctly, just as pageable (slower) copies.

Functions:

  • host_register –

    Register a host range as device-accessible pinned memory.

  • host_unregister –

    Release a host range previously registered by host_register.

host_register(ptr, num_bytes)

Register a host range as device-accessible pinned memory.

Parameters:

  • ptr

    (int) –

    Base address of the host range.

  • num_bytes

    (int) –

    Size of the host range in bytes.

Returns:

  • bool –

    True if the range was registered and later needs host_unregister.

Source code in vllm/v1/kv_offload/cpu/host_register.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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def host_register(ptr: int, num_bytes: int) -> bool:
    """Register a host range as device-accessible pinned memory.

    Args:
        ptr: Base address of the host range.
        num_bytes: Size of the host range in bytes.

    Returns:
        True if the range was registered and later needs `host_unregister`.

    """
    if not ptr or num_bytes <= 0:
        return False

    if current_platform.is_cuda_alike():
        cudart = _load_cudart()
        if cudart is None:
            return False
        result = cudart.cudaHostRegister(ptr, num_bytes)
        if result != 0:
            cudart.cudaGetLastError()
            logger.warning(
                "cudaHostRegister failed (code=%d) — transfers will still work "
                "but may be slower (unpinned DMA)",
                result,
            )
            return False
        return True

    if current_platform.is_xpu():
        if not torch.ops._C.xpu_host_register(ptr, num_bytes):
            logger.warning(
                "xpu_host_register failed — transfers will still work "
                "but may be slower (unpinned DMA)"
            )
            return False
        return True

    logger.info_once(
        "Host registration is not implemented on %s; KV offload transfers will "
        "use unpinned host memory.",
        current_platform.device_name,
    )
    return False

host_unregister(ptr)

Release a host range previously registered by host_register.

Source code in vllm/v1/kv_offload/cpu/host_register.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def host_unregister(ptr: int) -> None:
    """Release a host range previously registered by `host_register`."""
    if not ptr:
        return

    if current_platform.is_cuda_alike():
        cudart = _load_cudart()
        if cudart is None:
            return
        result = cudart.cudaHostUnregister(ptr)
        if result != 0:
            cudart.cudaGetLastError()
            logger.warning("cudaHostUnregister failed (code=%d)", result)
        return

    if current_platform.is_xpu():
        if not torch.ops._C.xpu_host_unregister(ptr):
            logger.warning("xpu_host_unregister failed")
        return