75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import torch
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.platforms import AttentionBackendEnum
|
|
from fastvideo.platforms.interface import (DeviceCapability, Platform, PlatformEnum)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class MpsPlatform(Platform):
|
|
_enum = PlatformEnum.MPS
|
|
device_name: str = "mps"
|
|
device_type: str = "mps"
|
|
dispatch_key: str = "MPS"
|
|
device_control_env_var: str = "MPS_VISIBLE_DEVICES"
|
|
|
|
@classmethod
|
|
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_device_name(cls, device_id: int = 0) -> str:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_device_uuid(cls, device_id: int = 0) -> str:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
|
|
if enforce_eager:
|
|
logger.warning("To see benefits of async output processing, enable MPS "
|
|
"graph. Since, enforce-eager is enabled, async output "
|
|
"processor cannot be used")
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> float:
|
|
return 0.0
|
|
|
|
@classmethod
|
|
def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
|
|
dtype: torch.dtype) -> str:
|
|
if selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
|
|
raise NotImplementedError("VIDEO_SPARSE_ATTN is not supported on MPS. Unset "
|
|
"FASTVIDEO_ATTENTION_BACKEND or set it to TORCH_SDPA.")
|
|
# MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible.
|
|
logger.info("Using Torch SDPA backend for MPS.")
|
|
return "fastvideo.attention.backends.sdpa.SDPABackend"
|
|
|
|
@classmethod
|
|
def get_device_communicator_cls(cls) -> str:
|
|
# Use base communicator for MPS
|
|
return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"
|
|
|
|
@classmethod
|
|
def seed_everything(cls, seed: int | None = None) -> None:
|
|
"""Set the seed for MPS device."""
|
|
if seed is not None:
|
|
import random
|
|
|
|
import numpy as np
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
# MPS doesn't have manual_seed_all like CUDA
|
|
# The manual_seed above should be sufficient
|