vllm/cacheflow/utils.py

21 lines
335 B
Python
Raw Normal View History

2023-02-09 19:26:50 +08:00
import enum
class Device(enum.Enum):
GPU = enum.auto()
CPU = enum.auto()
class Counter:
def __init__(self, start: int = 0) -> None:
self.counter = start
2023-02-14 09:19:27 +08:00
def __next__(self) -> int:
2023-02-09 19:26:50 +08:00
id = self.counter
self.counter += 1
return id
def reset(self) -> None:
self.counter = 0