| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Sequence |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | |
| 8 | class UvaBuffer: |
| 9 | def __init__(self, size: int | Sequence[int], dtype: torch.dtype): |
| 10 | self.cpu = torch.zeros(size, dtype=dtype, device="cpu") |
| 11 | self.np = self.cpu.numpy() |
| 12 | |
| 13 | def uva(self, n: int | None = None): |
| 14 | return self.cpu[:n] if n is not None else self.cpu |