forked from mooncake-track/Mooncake
108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
import ctypes
|
|
import logging
|
|
import os
|
|
import threading
|
|
from importlib import resources
|
|
from typing import Dict, Final
|
|
import torch_npu
|
|
|
|
from torch import device as torch_device
|
|
from torch_npu.npu.memory import NPUPluggableAllocator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class UBShmemAllocator:
|
|
_instances: Dict[torch_device, NPUPluggableAllocator] = {}
|
|
_lock: Final = threading.Lock()
|
|
_supports_fabric: bool = False
|
|
_probe_done: bool = False
|
|
|
|
@classmethod
|
|
def _get_so_path(cls) -> str:
|
|
"""Dynamically locate ubshmem_fabric_allocator.so in the mooncake package installation"""
|
|
try:
|
|
# Attempt to locate package resource
|
|
with resources.path("mooncake", "ubshmem_fabric_allocator.so") as so_path:
|
|
if so_path.exists():
|
|
return str(so_path)
|
|
except (ImportError, FileNotFoundError, TypeError):
|
|
pass
|
|
|
|
# Fallback strategy: check in package location via import metadata
|
|
try:
|
|
import mooncake
|
|
|
|
base_path = os.path.dirname(os.path.abspath(mooncake.__file__))
|
|
so_path = os.path.join(base_path, "ubshmem_fabric_allocator.so")
|
|
if os.path.exists(so_path):
|
|
return so_path
|
|
except (ImportError, FileNotFoundError, TypeError):
|
|
raise ImportError(
|
|
"UBShmemAllocator require mooncake-transfer-engine with USE_UBSHMEM enabled."
|
|
)
|
|
|
|
@classmethod
|
|
def _probe_fabric_memory_support(cls, so_path: str) -> bool:
|
|
"""
|
|
Probe whether the system supports fabric memory by calling a C++ function
|
|
that attempts aclrtMallocPhysical with fabric memory support.
|
|
The shared library exports a symbol like:
|
|
extern "C" bool mc_probe_ub_fabric_support(int device_id);
|
|
"""
|
|
try:
|
|
lib = ctypes.CDLL(so_path)
|
|
|
|
# Try to get the probe function
|
|
probe_func = lib.mc_probe_ub_fabric_support
|
|
probe_func.argtypes = [ctypes.c_int]
|
|
probe_func.restype = ctypes.c_bool
|
|
|
|
# Use device 0 for probing
|
|
dev_id = 0
|
|
supported = probe_func(dev_id)
|
|
if supported:
|
|
logger.info(f"Supports Fabric Memory with aclMallocPhysical")
|
|
else:
|
|
logger.info("Fabric memory not supported")
|
|
return supported
|
|
|
|
except AttributeError:
|
|
logger.warning(
|
|
"Symbol 'mc_probe_ub_fabric_support' not found in ubshmem_fabric_allocator.so. "
|
|
"Assuming fabric memory is NOT supported (you may need to update the library)."
|
|
)
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Failed to probe fabric memory support: {e}")
|
|
return False
|
|
|
|
@classmethod
|
|
def detect_mem_backend(cls) -> bool:
|
|
"""Public API: check if fabric memory is supported."""
|
|
if not cls._probe_done:
|
|
with cls._lock:
|
|
if cls._probe_done:
|
|
return cls._supports_fabric
|
|
so_path = None
|
|
try:
|
|
so_path = cls._get_so_path()
|
|
# First try dedicated probe function
|
|
cls._supports_fabric = cls._probe_fabric_memory_support(so_path)
|
|
except Exception as e:
|
|
logger.error(f"Critical error during fabric memory probe setup: {e}")
|
|
cls._supports_fabric = False
|
|
|
|
cls._probe_done = True
|
|
return cls._supports_fabric
|
|
|
|
@classmethod
|
|
def get_allocator(cls, device: torch_device) -> NPUPluggableAllocator:
|
|
with cls._lock:
|
|
if device not in cls._instances:
|
|
so_path = cls._get_so_path()
|
|
cls._instances[device] = NPUPluggableAllocator(
|
|
so_path, "mc_ub_fabric_malloc", "mc_ub_fabric_free"
|
|
)
|
|
return cls._instances[device]
|