vllm_mlx.optimizations

Hardware detection and system information for vllm-mlx.

View the complete module source at #L1-L209.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.optimizations

Hardware detection and system information for vllm-mlx.

This module provides: - Hardware detection for Apple Silicon (M1, M2, M3, M4 series) - System memory detection - Memory bandwidth benchmarking

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 vllm_mlx.optimizations import ( detect_hardware, get_optimization_status, benchmark_memory_bandwidth, )

vllm_mlx.optimizations.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.optimizations.HARDWARE_PROFILES module-attribute

HARDWARE_PROFILES = {'M1': {'bandwidth': 68.25, 'gpu_cores': 8}, 'M1 Pro': {'bandwidth': 200, 'gpu_cores': 16}, 'M1 Max': {'bandwidth': 400, 'gpu_cores': 32}, 'M1 Ultra': {'bandwidth': 800, 'gpu_cores': 64}, 'M2': {'bandwidth': 100, 'gpu_cores': 10}, 'M2 Pro': {'bandwidth': 200, 'gpu_cores': 19}, 'M2 Max': {'bandwidth': 400, 'gpu_cores': 38}, 'M2 Ultra': {'bandwidth': 800, 'gpu_cores': 76}, 'M3': {'bandwidth': 100, 'gpu_cores': 10}, 'M3 Pro': {'bandwidth': 150, 'gpu_cores': 18}, 'M3 Max': {'bandwidth': 400, 'gpu_cores': 40}, 'M3 Ultra': {'bandwidth': 800, 'gpu_cores': 80}, 'M4': {'bandwidth': 120, 'gpu_cores': 10}, 'M4 Pro': {'bandwidth': 273, 'gpu_cores': 20}, 'M4 Max': {'bandwidth': 546, 'gpu_cores': 40}, 'M4 Ultra': {'bandwidth': 800, 'gpu_cores': 80}}

vllm_mlx.optimizations.HardwareInfo dataclass

HardwareInfo(chip_name: str, total_memory_gb: float, memory_bandwidth_gbs: float, gpu_cores: int)

Hardware information for Apple Silicon.

vllm_mlx.optimizations.HardwareInfo.chip_name instance-attribute

chip_name: str

vllm_mlx.optimizations.HardwareInfo.total_memory_gb instance-attribute

total_memory_gb: float

vllm_mlx.optimizations.HardwareInfo.memory_bandwidth_gbs instance-attribute

memory_bandwidth_gbs: float

vllm_mlx.optimizations.HardwareInfo.gpu_cores instance-attribute

gpu_cores: int

vllm_mlx.optimizations.get_system_memory_gb

get_system_memory_gb() -> float

Get actual system memory in GB.

Returns:

• float –

  Total system memory in GB (unified memory on Apple Silicon)

Source code in vllm_mlx/optimizations.py

def get_system_memory_gb() -> float:
    """
    Get actual system memory in GB.

    Returns:
        Total system memory in GB (unified memory on Apple Silicon)
    """
    try:
        import subprocess

        result = subprocess.run(
            ["sysctl", "-n", "hw.memsize"],
            capture_output=True,
            text=True,
            check=True,
        )
        mem_bytes = int(result.stdout.strip())
        return mem_bytes / (1024**3)
    except Exception:
        # Fallback: try to get from MLX device info
        try:
            device_info = mx.device_info()
            if "memory_size" in device_info:
                return device_info["memory_size"] / (1024**3)
        except Exception:
            pass
        return 16.0  # Conservative default

vllm_mlx.optimizations.detect_hardware

detect_hardware() -> HardwareInfo

Detect Apple Silicon hardware and return info.

Memory is detected dynamically from the system. Other specs (bandwidth, GPU cores) come from known chip profiles.

Returns:

• HardwareInfo –

  HardwareInfo with detected hardware specifications

Source code in vllm_mlx/optimizations.py

def detect_hardware() -> HardwareInfo:
    """
    Detect Apple Silicon hardware and return info.

    Memory is detected dynamically from the system.
    Other specs (bandwidth, GPU cores) come from known chip profiles.

    Returns:
        HardwareInfo with detected hardware specifications
    """
    try:
        device_info = mx.device_info()
        device_name = device_info.get("device_name", "")
        actual_memory_gb = get_system_memory_gb()

        # Match with known profiles (check longest names first)
        sorted_profiles = sorted(
            HARDWARE_PROFILES.items(), key=lambda x: len(x[0]), reverse=True
        )

        for chip_name, profile in sorted_profiles:
            if chip_name in device_name:
                return HardwareInfo(
                    chip_name=chip_name,
                    total_memory_gb=actual_memory_gb,
                    memory_bandwidth_gbs=profile["bandwidth"],
                    gpu_cores=profile["gpu_cores"],
                )

        # Unknown chip
        return HardwareInfo(
            chip_name="Unknown",
            total_memory_gb=actual_memory_gb,
            memory_bandwidth_gbs=200,
            gpu_cores=16,
        )

    except Exception as e:
        logger.warning(f"Failed to detect hardware: {e}")
        return HardwareInfo(
            chip_name="Unknown",
            total_memory_gb=get_system_memory_gb(),
            memory_bandwidth_gbs=200,
            gpu_cores=16,
        )

vllm_mlx.optimizations.benchmark_memory_bandwidth

benchmark_memory_bandwidth() -> dict

Benchmark actual memory bandwidth achieved.

Returns:

• dict –

  dict with bandwidth measurements for different array sizes

Source code in vllm_mlx/optimizations.py

def benchmark_memory_bandwidth() -> dict:
    """
    Benchmark actual memory bandwidth achieved.

    Returns:
        dict with bandwidth measurements for different array sizes
    """
    import time

    sizes_mb = [1, 4, 16]
    results = {}

    for size_mb in sizes_mb:
        elements = (size_mb * 1024 * 1024) // 4  # float32
        a = mx.random.normal((elements,))
        b = mx.random.normal((elements,))
        mx.eval(a, b)

        start = time.perf_counter()
        for _ in range(10):
            c = a + b
            mx.eval(c)
        elapsed = time.perf_counter() - start

        # read a, read b, write c = 3x size
        total_bytes = 3 * size_mb * 1024 * 1024 * 10
        bandwidth_gbs = (total_bytes / elapsed) / 1e9

        results[f"{size_mb}MB"] = f"{bandwidth_gbs:.1f} GB/s"

    return results

vllm_mlx.optimizations.get_optimization_status

get_optimization_status() -> dict

Get current hardware and MLX status.

Returns:

• dict –

  dict with hardware info and MLX configuration

Source code in vllm_mlx/optimizations.py

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,
            "memory_bandwidth_gbs": hw.memory_bandwidth_gbs,
            "gpu_cores": hw.gpu_cores,
            "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",
        },
    }

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.optimizations.HardwareInfo · class

vllm_mlx.optimizations.HardwareInfo(chip_name: str, total_memory_gb: float, memory_bandwidth_gbs: float, gpu_cores: int)

Hardware information for Apple Silicon.

Parameters

Name	Type	Required	Default	Description
chip_name	str	yes	none	Required constructor field.
total_memory_gb	float	yes	none	Required constructor field.
memory_bandwidth_gbs	float	yes	none	Required constructor field.
gpu_cores	int	yes	none	Required constructor field.

Returns

• Constructs: vllm_mlx.optimizations.HardwareInfo

Exceptions and behavior

Class HardwareInfo declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L34-L40.

vllm_mlx.optimizations.get_system_memory_gb · function

vllm_mlx.optimizations.get_system_memory_gb() -> float

Get actual system memory in GB.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: mem_bytes / 1024 ** 3; device_info['memory_size'] / 1024 ** 3; 16.0

Exceptions and behavior

Function get_system_memory_gb calls subprocess.run, int, result.stdout.strip, mx.device_info; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L68-L94.

vllm_mlx.optimizations.detect_hardware · function

vllm_mlx.optimizations.detect_hardware() -> HardwareInfo

Detect Apple Silicon hardware and return info.

Parameters

This callable has no explicit inputs.

Returns

• Type: HardwareInfo
• Direct return expressions: HardwareInfo(chip_name=chip_name, total_memory_gb=actual_memory_gb, memory_bandwidth_gbs=profile['bandwidth'], gpu_core…; HardwareInfo(chip_name='Unknown', total_memory_gb=actual_memory_gb, memory_bandwidth_gbs=200, gpu_cores=16); HardwareInfo(chip_name='Unknown', total_memory_gb=get_system_memory_gb(), memory_bandwidth_gbs=200, gpu_cores=16)

Exceptions and behavior

Function detect_hardware calls mx.device_info, device_info.get, get_system_memory_gb, sorted; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L97-L141.

vllm_mlx.optimizations.benchmark_memory_bandwidth · function

vllm_mlx.optimizations.benchmark_memory_bandwidth() -> dict

Benchmark actual memory bandwidth achieved.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: results

Exceptions and behavior

Function benchmark_memory_bandwidth calls mx.random.normal, mx.eval, time.perf_counter, range; returns results. No direct raise statement appears in this definition.

View source #L144-L174.

vllm_mlx.optimizations.get_optimization_status · function

vllm_mlx.optimizations.get_optimization_status() -> dict

Get current hardware and MLX status.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: {'hardware': {'chip': hw.chip_name, 'total_memory_gb': hw.total_memory_gb, 'memory_bandwidth_gbs': hw.memory_bandwidth_…

Exceptions and behavior

Function get_optimization_status calls detect_hardware, mx.device_info, hasattr, device_info.get; returns {'hardware': {'chip': hw.chip_name, 'total_memory_gb': hw.total_memory_gb, 'memory_bandwidth_gbs': hw.memory_bandwidth_…. No direct raise statement appears in this definition.

View source #L177-L209.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
HardwareInfo	class	HardwareInfo(chip_name: str, total_memory_gb: float, memory_bandwidth_gbs: float, gpu_cores: int)	Hardware information for Apple Silicon.	#L34-L40
get_system_memory_gb	function	get_system_memory_gb() -> float	Get actual system memory in GB.	#L68-L94
detect_hardware	function	detect_hardware() -> HardwareInfo	Detect Apple Silicon hardware and return info.	#L97-L141
benchmark_memory_bandwidth	function	benchmark_memory_bandwidth() -> dict	Benchmark actual memory bandwidth achieved.	#L144-L174
get_optimization_status	function	get_optimization_status() -> dict	Get current hardware and MLX status.	#L177-L209