platforms ¶

@classmethod
def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
                         dtype: torch.dtype) -> str:
    """Get the attention backend class of a device."""
    return ""

@classmethod
def get_cpu_architecture(cls) -> CpuArchEnum:
    """Get the CPU architecture of the current platform."""
    return CpuArchEnum.UNSPECIFIED

@classmethod
def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> float:
    """
    Return the memory usage in bytes.
    """
    raise NotImplementedError

@classmethod
def get_device_capability(
    cls,
    device_id: int = 0,
) -> DeviceCapability | None:
    """Stateless version of :func:`torch.cuda.get_device_capability`."""
    return None

@classmethod
def get_device_communicator_cls(cls) -> str:
    """
    Get device specific communicator class for distributed communication.
    """
    return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"  # noqa

@classmethod
def get_device_name(cls, device_id: int = 0) -> str:
    """Get the name of a device."""
    raise NotImplementedError

@classmethod
def get_device_total_memory(cls, device_id: int = 0) -> int:
    """Get the total memory of a device in bytes."""
    raise NotImplementedError

@classmethod
def get_device_uuid(cls, device_id: int = 0) -> str:
    """Get the uuid of a device, e.g. the PCI bus ID."""
    raise NotImplementedError

@classmethod
def get_torch_device(cls):
    """
    Check if the current platform supports torch device.
    """
    raise NotImplementedError

@classmethod
def has_device_capability(
    cls,
    capability: tuple[int, int] | int,
    device_id: int = 0,
) -> bool:
    """
    Test whether this platform is compatible with a device capability.

    The ``capability`` argument can either be:

    - A tuple ``(major, minor)``.
    - An integer ``<major><minor>``. (See :meth:`DeviceCapability.to_int`)
    """
    current_capability = cls.get_device_capability(device_id=device_id)
    if current_capability is None:
        return False

    if isinstance(capability, tuple):
        return current_capability >= capability

    return current_capability.to_int() >= capability

@classmethod
def inference_mode(cls):
    """A device-specific wrapper of `torch.inference_mode`.

    This wrapper is recommended because some hardware backends such as TPU
    do not support `torch.inference_mode`. In such a case, they will fall
    back to `torch.no_grad` by overriding this method.
    """
    return torch.inference_mode(mode=True)

@classmethod
def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
    """
    Check if the current platform supports async output.
    """
    raise NotImplementedError

def is_cuda_alike(self) -> bool:
    """Stateless version of :func:`torch.cuda.is_available`."""
    return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)

@classmethod
def seed_everything(cls, seed: int | None = None) -> None:
    """
    Set the seed of each random module.
    `torch.manual_seed` will set seed on all devices.

    Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
    """
    if seed is not None:
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

@classmethod
def verify_model_arch(cls, model_arch: str) -> None:
    """
    Verify whether the current platform supports the specified model
    architecture.

    - This will raise an Error or Warning based on the model support on
    the current platform.
    - By default all models are considered supported.
    """
    pass

@classmethod
def verify_quantization(cls, quant: str) -> None:
    """
    Verify whether the quantization is supported by the current platform.
    """
    if cls.supported_quantization and \
        quant not in cls.supported_quantization:
        raise ValueError(f"{quant} quantization is currently not supported in "
                         f"{cls.device_name}.")