93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Per-device byte accounting for session state (RFC #4480).
|
|
|
|
A session holds bytes in more than one place -- device tensors, host tensors,
|
|
NumPy frame buffers -- and those are different resources: a GPU pool that is
|
|
scarce and contended, and host memory that is not. Summing them into a single
|
|
number answers no question anyone has, so the accounting here is keyed by the
|
|
device the bytes live on and is never reduced to a scalar by this module.
|
|
|
|
Two details make the numbers honest:
|
|
|
|
* **Storages, not elements.** ``numel() * element_size()`` measures the logical
|
|
view, but a narrow slice of a large buffer keeps the whole allocation alive.
|
|
Counting ``untyped_storage().nbytes()`` reports what is actually pinned.
|
|
* **Views deduplicated.** Two tensors sharing one storage are counted once, via
|
|
the caller-supplied ``seen`` set. Pass the *same* set across every value in a
|
|
session so an alias between two state objects (or between an object and an
|
|
attribute) does not double-count.
|
|
|
|
Anything that is not a tensor, an array, or a container of those contributes
|
|
nothing. That is not a heuristic: in this codebase the only device memory a
|
|
session can hold is a ``torch.Tensor``. CUDA graph pools, collective buffers
|
|
and the paged KV blocks of PR #4534 are owned by the engine, never by a session.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
# ``(device string, storage address)`` of a storage already counted.
|
|
SeenStorages = set[tuple[str, int]]
|
|
|
|
_HOST = "cpu"
|
|
|
|
|
|
def device_bytes(value: object, seen: SeenStorages) -> dict[str, int]:
|
|
"""Bytes pinned by ``value``, keyed by the device they live on.
|
|
|
|
``seen`` accumulates the storages already counted and is mutated in place;
|
|
share one set across a whole session to deduplicate aliases.
|
|
"""
|
|
totals: dict[str, int] = {}
|
|
_accumulate(value, seen, totals)
|
|
return totals
|
|
|
|
|
|
def merge_bytes(into: dict[str, int], other: Mapping[str, int]) -> dict[str, int]:
|
|
"""Add ``other`` into ``into`` per device, in place."""
|
|
for device, count in other.items():
|
|
into[device] = into.get(device, 0) + count
|
|
return into
|
|
|
|
|
|
def _accumulate(value: object, seen: SeenStorages, totals: dict[str, int]) -> None:
|
|
if isinstance(value, torch.Tensor):
|
|
storage = value.untyped_storage()
|
|
device = str(value.device)
|
|
key = (device, storage.data_ptr())
|
|
if key in seen:
|
|
return
|
|
seen.add(key)
|
|
totals[device] = totals.get(device, 0) + storage.nbytes()
|
|
return
|
|
|
|
if isinstance(value, np.ndarray):
|
|
# ``base is not None`` marks a view; NumPy views are cheap to alias and
|
|
# carry no address we can key on as reliably as a torch storage, so the
|
|
# owning array is what gets counted.
|
|
owner = value if value.base is None else value.base
|
|
if isinstance(owner, np.ndarray):
|
|
key = (_HOST, owner.__array_interface__["data"][0])
|
|
if key in seen:
|
|
return
|
|
seen.add(key)
|
|
totals[_HOST] = totals.get(_HOST, 0) + int(owner.nbytes)
|
|
return
|
|
|
|
if isinstance(value, Mapping):
|
|
for item in value.values():
|
|
_accumulate(item, seen, totals)
|
|
return
|
|
|
|
# ``str``/``bytes`` are Sequences but hold no device memory; excluding them
|
|
# also stops the recursion from walking a string one character at a time.
|
|
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
|
|
for item in value:
|
|
_accumulate(item, seen, totals)
|
|
return
|