File size: 1,246 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """PF1-CAUSE helper (mounted as vllm/v1/worker/gpu/dsv41_prefill_cause.py): host-only allocator counters for a span row.
Only while a PF1-R2 trace chunk is open (prefill steps of an ON arm); otherwise an empty dict and no torch call. Reads
the caching allocator's own statistics (no device read, no synchronization, no allocation). Numbers only.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING: # the same module as reviewed in this checkout (offload/prefill_trace/runtime.py)
from offload.prefill_trace import runtime as _pf
else:
from vllm.v1.worker.gpu import dsv41_prefill_trace as _pf
Scalar = str | int | float | bool | None
STATS = {'alloc_retries': 'num_alloc_retries', 'device_allocs': 'num_device_alloc', 'device_frees': 'num_device_free',
'reserved_bytes': 'reserved_bytes.all.current', 'active_bytes': 'active_bytes.all.current',
'segments': 'segment.all.current'}
def alloc(**extra: Scalar) -> dict[str, Scalar]:
rec = _pf._recorder
if not _pf.ENABLED or rec is None or rec.context is None:
return {}
import torch
stats = torch.cuda.memory_stats()
return {**extra, **{name: int(stats.get(key, -1)) for name, key in STATS.items()}}
|