wehub-resource-sync
76 lines
1.9 KiB
Python
76 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Hardware detection and system information for oMLX.
|
|
|
|
This module provides:
|
|
- Hardware detection for Apple Silicon
|
|
- System memory detection
|
|
|
|
Note: mlx-lm already includes optimized implementations internally:
|
|
- Flash Attention via mx.fast.scaled_dot_product_attention
|
|
- Efficient memory management
|
|
- Optimized Metal kernels
|
|
|
|
No additional optimization is needed - mlx-lm is already fast out of the box.
|
|
|
|
Usage:
|
|
from omlx.optimizations import (
|
|
detect_hardware,
|
|
get_optimization_status,
|
|
)
|
|
"""
|
|
|
|
import logging
|
|
|
|
import mlx.core as mx
|
|
|
|
# Re-export from hardware module for backward compatibility
|
|
from omlx.utils.hardware import (
|
|
HardwareInfo,
|
|
detect_hardware,
|
|
get_total_memory_gb as get_system_memory_gb,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
__all__ = [
|
|
"HardwareInfo",
|
|
"detect_hardware",
|
|
"get_system_memory_gb",
|
|
"get_optimization_status",
|
|
]
|
|
|
|
|
|
|
|
def get_optimization_status() -> dict:
|
|
"""
|
|
Get current hardware and MLX status.
|
|
|
|
Returns:
|
|
dict with hardware info and MLX configuration
|
|
"""
|
|
hw = detect_hardware()
|
|
device_info = mx.device_info()
|
|
flash_available = hasattr(mx, "fast") and hasattr(
|
|
mx.fast, "scaled_dot_product_attention"
|
|
)
|
|
|
|
return {
|
|
"hardware": {
|
|
"chip": hw.chip_name,
|
|
"total_memory_gb": hw.total_memory_gb,
|
|
"device_name": device_info.get("device_name", "Unknown"),
|
|
},
|
|
"mlx_memory": {
|
|
"active_bytes": mx.get_active_memory(),
|
|
"cache_bytes": mx.get_cache_memory(),
|
|
"peak_bytes": mx.get_peak_memory(),
|
|
},
|
|
"mlx_lm_features": {
|
|
"flash_attention": "built-in" if flash_available else "not available",
|
|
"metal_kernels": "optimized for Apple Silicon",
|
|
"kv_cache": "managed by mlx-lm",
|
|
"quantization": "4-bit and 8-bit supported",
|
|
},
|
|
}
|