diff --git a/vllm_metal/model_runner.py b/vllm_metal/model_runner.py index 7ba5d10e..1570ff27 100644 --- a/vllm_metal/model_runner.py +++ b/vllm_metal/model_runner.py @@ -11,6 +11,7 @@ from vllm_metal.config import get_config from vllm_metal.mlx_backend.cache import PagedKVCache +from vllm_metal.platform import set_wired_limit if TYPE_CHECKING: from vllm.config import VllmConfig @@ -44,6 +45,7 @@ def load_model(self) -> None: model_name = model_config.model logger.info(f"Loading model: {model_name}") + set_wired_limit() # Load model and tokenizer using mlx_lm self.model, self.tokenizer = mlx_load( diff --git a/vllm_metal/platform.py b/vllm_metal/platform.py index 4b55f07c..400ce199 100644 --- a/vllm_metal/platform.py +++ b/vllm_metal/platform.py @@ -19,6 +19,29 @@ logger = logging.getLogger(__name__) +def set_wired_limit() -> None: + """Set Metal wired memory limit for optimal GPU performance. + + Pins model weights in GPU-accessible memory to prevent memory paging + and GPU stalls during inference. + + See: https://github.com/ml-explore/mlx-lm/pull/652 + """ + try: + import mlx.core as mx + + device_info = mx.metal.device_info() + max_wired = device_info.get("max_recommended_working_set_size", 0) + if max_wired > 0: + if hasattr(mx, "set_wired_limit"): + mx.set_wired_limit(max_wired) + elif hasattr(mx.metal, "set_wired_limit"): + mx.metal.set_wired_limit(max_wired) + logger.info(f"Set Metal wired_limit to {max_wired / (1024**3):.1f} GB") + except Exception as e: + logger.warning(f"Failed to set wired_limit: {e}") + + class MetalPlatform(Platform): """Platform implementation for Apple Silicon Metal/MLX. diff --git a/vllm_metal/v1/worker.py b/vllm_metal/v1/worker.py index 454ba6a5..cff92209 100644 --- a/vllm_metal/v1/worker.py +++ b/vllm_metal/v1/worker.py @@ -22,7 +22,7 @@ from vllm.v1.worker.worker_base import WorkerBase from vllm_metal.config import get_config -from vllm_metal.platform import MetalPlatform +from vllm_metal.platform import MetalPlatform, set_wired_limit if TYPE_CHECKING: from vllm_metal.v1.model_runner import MetalModelRunner @@ -95,6 +95,7 @@ def init_device(self) -> None: ) mx.set_default_device(mx.Device(device_type)) logger.info(f"MLX device set to: {mx.default_device()}") + set_wired_limit() # Use MetalPlatform.get_torch_device() to properly support MPS when available. # This ensures consistency with the platform's device selection logic and