vllm/cacheflow/utils.py
2023-02-14 01:19:27 +00:00

21 lines
335 B
Python

import enum
class Device(enum.Enum):
GPU = enum.auto()
CPU = enum.auto()
class Counter:
def __init__(self, start: int = 0) -> None:
self.counter = start
def __next__(self) -> int:
id = self.counter
self.counter += 1
return id
def reset(self) -> None:
self.counter = 0