Skip to content

vllm.v1.simple_kv_offload.host_buffer

One host mapping backing a rank's CPU KV caches.

Functions:

alloc_host_buffer(nbytes, hugepage_size=None)

Map and populate an int8 CPU tensor of whole backing pages.

HugeTLB pages never fall back to ordinary pages.

Source code in vllm/v1/simple_kv_offload/host_buffer.py 
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
def alloc_host_buffer(nbytes: int, hugepage_size: int | None = None) -> torch.Tensor:
    """Map and populate an int8 CPU tensor of whole backing pages.

    HugeTLB pages never fall back to ordinary pages.
    """
    page_size = hugepage_size or mmap.PAGESIZE
    nbytes = round_up(max(nbytes, 1), page_size)
    kind = f"{page_size >> 20} MiB HugeTLB" if hugepage_size else "ordinary"
    flags = mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS
    if hugepage_size:
        flags |= _MAP_HUGETLB | (hugepage_size.bit_length() - 1) << _MAP_HUGE_SHIFT
    try:
        buf = mmap.mmap(-1, nbytes, flags=flags)
    except OSError as e:
        raise RuntimeError(
            f"Cannot map {nbytes / 2**30:.2f} GiB of {kind} pages: {e}. "
            "Reserve the pages on the host and request them for the container."
        ) from e
    # Keep pinned pages out of fork() children.
    buf.madvise(mmap.MADV_DONTFORK)
    if not hugepage_size and os.environ.get("THP_MEM_ALLOC_ENABLE") == "1":
        try:
            buf.madvise(mmap.MADV_HUGEPAGE)
        except OSError as e:
            logger.warning("MADV_HUGEPAGE failed: %s", e)
    populate = _madvise_populate_write if hugepage_size else _get_populate_write_fn(buf)
    try:
        # A pool or cgroup shortfall fails here instead of as SIGBUS or an
        # OOM kill while serving.
        populate(buf, 0, nbytes)
    except OSError as e:
        raise RuntimeError(
            f"Cannot populate {nbytes / 2**30:.2f} GiB of {kind} pages: "
            f"{e}. The container's memory or hugetlb limit is too small."
        ) from e
    return torch.frombuffer(buf, dtype=torch.int8)

parse_hugepage_size(value)

Parse cpu_hugepage_block_size; None or "none" disables it.

Source code in vllm/v1/simple_kv_offload/host_buffer.py 
25
26
27
28
29
30
31
32
33
34
35
def parse_hugepage_size(value: object) -> int | None:
    """Parse ``cpu_hugepage_block_size``; ``None`` or ``"none"`` disables it."""
    key = str(value or "none").upper()
    if key == "NONE":
        return None
    if key not in HUGEPAGE_SIZES:
        raise ValueError(
            f"Unsupported cpu_hugepage_block_size {value!r}; "
            f"expected one of {sorted(HUGEPAGE_SIZES)} or none"
        )
    return HUGEPAGE_SIZES[key]