Set wired_limit to fix low GPU utilization

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(

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.
 

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