feat: Add --gpu-memory-utilization, repetition_penalty, fix OOM on large models

- Add --gpu-memory-utilization CLI flag (default 0.90) to control Metal
  soft allocation limit and emergency cache clear threshold
- Fix OOM SIGKILL on large models (200GB+): clear traceback references
  between strict=True/False retry in load_model_with_fallback() to free
  memory from the failed first load before retrying
- Add repetition_penalty/frequency_penalty/presence_penalty support
  for chat and completion endpoints

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: janhilgard

vllm_mlx/api/models.py

@@ -172,6 +172,10 @@ class ChatCompletionRequest(BaseModel):
     # MLLM-specific parameters
     video_fps: float | None = None
     video_max_frames: int | None = None
+    # Sampling penalties
+    repetition_penalty: float | None = None  # mlx-lm style (>1.0 penalizes)
+    frequency_penalty: float | None = None  # OpenAI style (0-2)
+    presence_penalty: float | None = None  # OpenAI style (0-2)
     # Request timeout in seconds (None = use server default)
     timeout: float | None = None
 
@@ -235,6 +239,10 @@ class CompletionRequest(BaseModel):
     max_tokens: int | None = None
     stream: bool = False
     stop: list[str] | None = None
+    # Sampling penalties
+    repetition_penalty: float | None = None  # mlx-lm style (>1.0 penalizes)
+    frequency_penalty: float | None = None  # OpenAI style (0-2)
+    presence_penalty: float | None = None  # OpenAI style (0-2)
     # Request timeout in seconds (None = use server default)
     timeout: float | None = None
 

vllm_mlx/cli.py

@@ -37,6 +37,13 @@ def serve_command(args):
         print("Example: --enable-auto-tool-choice --tool-call-parser mistral")
         sys.exit(1)
 
+    # Validate gpu-memory-utilization range
+    if not (0.0 < args.gpu_memory_utilization <= 1.0):
+        print(
+            "Error: --gpu-memory-utilization must be between 0.0 (exclusive) and 1.0 (inclusive)"
+        )
+        sys.exit(1)
+
     # Configure server security settings
     server._api_key = args.api_key
     server._default_timeout = args.timeout
@@ -186,6 +193,7 @@ def serve_command(args):
         scheduler_config=scheduler_config,
         stream_interval=args.stream_interval if args.continuous_batching else 1,
         max_tokens=args.max_tokens,
+        gpu_memory_utilization=args.gpu_memory_utilization,
     )
 
     # Start server
@@ -680,6 +688,14 @@ def main():
         action="store_true",
         help="Enable continuous batching for multiple concurrent users (slower for single user)",
     )
+    serve_parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.90,
+        help="Fraction of device memory for Metal allocation limit and emergency "
+        "cache clear threshold (0.0-1.0, default: 0.90). Increase to 0.95 for "
+        "large models (200GB+) that need more memory headroom.",
+    )
     # Paged cache options (experimental)
     serve_parser.add_argument(
         "--use-paged-cache",

vllm_mlx/engine/batched.py

@@ -137,6 +137,7 @@ def __init__(
         scheduler_config: Any | None = None,
         stream_interval: int = 1,
         force_mllm: bool = False,
+        gpu_memory_utilization: float = 0.90,
     ):
         """
         Initialize the batched engine.
@@ -147,11 +148,14 @@ def __init__(
             scheduler_config: Optional scheduler configuration
             stream_interval: Tokens to batch before streaming (1=every token)
             force_mllm: Force loading as MLLM even if not auto-detected
+            gpu_memory_utilization: Fraction of device memory for Metal allocation
+                limit and emergency threshold (0.0-1.0, default 0.90)
         """
         self._model_name = model_name
         self._trust_remote_code = trust_remote_code
         self._scheduler_config = scheduler_config
         self._stream_interval = stream_interval
+        self._gpu_memory_utilization = gpu_memory_utilization
         self._is_mllm = force_mllm or is_mllm_model(model_name)
 
         self._model = None
@@ -283,13 +287,14 @@ async def _start_llm(self) -> None:
                     device_info.get("memory_size", 0),
                 )
                 if max_recommended > 0:
-                    soft_limit = int(max_recommended * 0.95)
+                    soft_limit = int(max_recommended * self._gpu_memory_utilization)
                     mx.set_memory_limit(soft_limit)
                     mx.set_cache_limit(32 * 1024 * 1024 * 1024)  # 32GB
+                    pct = self._gpu_memory_utilization * 100
                     logger.info(
                         f"Metal memory limits set: "
                         f"allocation_limit={soft_limit / 1e9:.1f}GB "
-                        f"(90% of {max_recommended / 1e9:.1f}GB), "
+                        f"({pct:.0f}% of {max_recommended / 1e9:.1f}GB), "
                         f"cache_limit=32GB"
                     )
         except Exception as e:
@@ -301,6 +306,7 @@ async def _start_llm(self) -> None:
             model_name=self._model_name,
             scheduler_config=scheduler_config,
             stream_interval=self._stream_interval,
+            gpu_memory_utilization=self._gpu_memory_utilization,
         )
 
         # Create async engine

vllm_mlx/engine_core.py

@@ -36,6 +36,7 @@ class EngineConfig:
     scheduler_config: Optional[SchedulerConfig] = None
     step_interval: float = 0.001  # 1ms between steps
     stream_interval: int = 1  # Tokens to batch before streaming (1=every token)
+    gpu_memory_utilization: float = 0.90  # Fraction of device memory for allocation
 
 
 class EngineCore:
@@ -150,8 +151,17 @@ async def _engine_loop(self) -> None:
         stream_interval = self.config.stream_interval
         use_simple_streaming = stream_interval == 1
 
-        # Emergency memory pressure threshold (245GB — raised for large models)
-        _memory_pressure_threshold = 245 * 1024 * 1024 * 1024
+        # Emergency memory pressure threshold — dynamic based on gpu_memory_utilization
+        _gpu_mem_util = self.config.gpu_memory_utilization
+        try:
+            _device_mem = mx.device_info().get(
+                "memory_size", 200 * 1024 * 1024 * 1024
+            )
+            _memory_pressure_threshold = int(
+                _device_mem * min(_gpu_mem_util + 0.05, 0.99)
+            )
+        except Exception:
+            _memory_pressure_threshold = 200 * 1024 * 1024 * 1024
         _memory_check_interval = 64
 
         while self._running:

vllm_mlx/scheduler.py

@@ -19,7 +19,7 @@
 
 import mlx.core as mx
 from mlx_lm.generate import BatchGenerator
-from mlx_lm.sample_utils import make_sampler
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
 
 from .memory_cache import MemoryAwarePrefixCache, MemoryCacheConfig
 from .paged_cache import PagedCacheManager
@@ -403,7 +403,7 @@ def _chunked_next(self=batch_gen):  # noqa: C901
 
                     if not is_cached:
                         padded = _left_pad_prompts(inputs_raw, max_length=max_length)
-                        prompt_cache = _make_cache(self.model, padding)
+                        prompt_cache = _make_cache(self.model, padding, self.max_kv_size)
                     else:
                         last_inputs = mx.array([p[-1:] for p in inputs_raw])
                         padded = _right_pad_prompts(inputs_raw, max_length=max_length)
@@ -644,6 +644,10 @@ def _mtp_step(
 
         # --- Apply logits processors + sample primary ---
         if any(logits_processors):
+            logger.debug(
+                f"[logits_proc] applying {sum(len(lp) for lp in logits_processors)} "
+                f"processors to batch_size={batch_size}"
+            )
             processed_logits = []
             for e in range(batch_size):
                 sample_logits = logits[e : e + 1]
@@ -1760,15 +1764,30 @@ def _schedule_waiting(self) -> List[Request]:
                 request.remaining_tokens = request.prompt_token_ids
                 tokens_to_process = request.prompt_token_ids
 
+            # Build per-request logits_processors from repetition_penalty
+            rep_penalty = request.sampling_params.repetition_penalty
+            lp = None
+            if rep_penalty and rep_penalty != 1.0:
+                lp = make_logits_processors(repetition_penalty=rep_penalty)
+                logger.info(
+                    f"[rep_penalty] request={request.request_id[:12]} "
+                    f"penalty={rep_penalty} processors={len(lp)}"
+                )
+
             # Insert into BatchGenerator with optional cache.
             # Wrap in try/except: if cache shapes are incompatible
             # (e.g. stale entry after BatchGenerator recreation),
             # fall back to no-cache insert instead of crashing.
+            insert_kwargs = {
+                "max_tokens": [request.sampling_params.max_tokens],
+                "caches": [cache_to_use] if cache_to_use else None,
+            }
+            if lp:
+                insert_kwargs["logits_processors"] = [lp]
             try:
                 uids = self.batch_generator.insert(
                     [tokens_to_process],
-                    max_tokens=[request.sampling_params.max_tokens],
-                    caches=[cache_to_use] if cache_to_use else None,
+                    **insert_kwargs,
                 )
             except Exception as e:
                 if cache_to_use is not None:
@@ -1781,10 +1800,10 @@ def _schedule_waiting(self) -> List[Request]:
                     request.cached_tokens = 0
                     request.remaining_tokens = request.prompt_token_ids
                     tokens_to_process = request.prompt_token_ids
+                    insert_kwargs["caches"] = None
                     uids = self.batch_generator.insert(
                         [tokens_to_process],
-                        max_tokens=[request.sampling_params.max_tokens],
-                        caches=None,
+                        **insert_kwargs,
                     )
                 else:
                     raise
@@ -1805,11 +1824,16 @@ def _schedule_waiting(self) -> List[Request]:
                     else ""
                 )
                 tokens_to_prefill = len(tokens_to_process)
+                rep_info = (
+                    f" rep_penalty={rep_penalty}"
+                    if rep_penalty and rep_penalty != 1.0
+                    else ""
+                )
                 logger.info(
                     f"[schedule] request={request.request_id[:12]} uid={uid} "
                     f"prompt_tokens={request.num_prompt_tokens} "
                     f"tokens_to_prefill={tokens_to_prefill}{cache_info} "
-                    f"max_tokens={request.sampling_params.max_tokens} "
+                    f"max_tokens={request.sampling_params.max_tokens}{rep_info} "
                     f"running={len(self.running)} waiting={len(self.waiting)}"
                 )
 

vllm_mlx/server.py

@@ -517,6 +517,7 @@ def load_model(
     stream_interval: int = 1,
     max_tokens: int = 32768,
     force_mllm: bool = False,
+    gpu_memory_utilization: float = 0.90,
 ):
     """
     Load a model (auto-detects MLLM vs LLM).
@@ -546,6 +547,7 @@ def load_model(
             scheduler_config=scheduler_config,
             stream_interval=stream_interval,
             force_mllm=force_mllm,
+            gpu_memory_utilization=gpu_memory_utilization,
         )
         # BatchedEngine will be started in lifespan (uvicorn's event loop)
         # Just log for now
@@ -1231,10 +1233,22 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
         f"prompt_chars={prompt_len} prompt_preview={prompt_preview!r}"
     )
 
+    # Resolve repetition penalty for completions
+    comp_rep_penalty = request.repetition_penalty
+    if comp_rep_penalty is None and request.frequency_penalty:
+        comp_rep_penalty = 1.0 + request.frequency_penalty
+    if comp_rep_penalty is None and request.presence_penalty:
+        comp_rep_penalty = 1.0 + request.presence_penalty
+
     if request.stream:
         return StreamingResponse(
             _disconnect_guard(
-                stream_completion(engine, prompts[0], request),
+                stream_completion(
+                    engine,
+                    prompts[0],
+                    request,
+                    repetition_penalty=comp_rep_penalty,
+                ),
                 raw_request,
             ),
             media_type="text/event-stream",
@@ -1248,14 +1262,16 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
     total_prompt_tokens = 0
 
     for i, prompt in enumerate(prompts):
+        gen_kwargs = {
+            "max_tokens": request.max_tokens or _default_max_tokens,
+            "temperature": _resolve_temperature(request.temperature),
+            "top_p": _resolve_top_p(request.top_p),
+            "stop": request.stop,
+        }
+        if comp_rep_penalty is not None:
+            gen_kwargs["repetition_penalty"] = comp_rep_penalty
         output = await _wait_with_disconnect(
-            engine.generate(
-                prompt=prompt,
-                max_tokens=request.max_tokens or _default_max_tokens,
-                temperature=_resolve_temperature(request.temperature),
-                top_p=_resolve_top_p(request.top_p),
-                stop=request.stop,
-            ),
+            engine.generate(prompt=prompt, **gen_kwargs),
             raw_request,
             timeout=timeout,
         )
@@ -1387,12 +1403,21 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
             # Inject JSON instruction into messages
             messages = _inject_json_instruction(messages, json_instruction)
 
+    # Resolve repetition penalty: explicit > frequency_penalty > presence_penalty
+    rep_penalty = request.repetition_penalty
+    if rep_penalty is None and request.frequency_penalty:
+        rep_penalty = 1.0 + request.frequency_penalty
+    if rep_penalty is None and request.presence_penalty:
+        rep_penalty = 1.0 + request.presence_penalty
+
     # Prepare kwargs
     chat_kwargs = {
         "max_tokens": request.max_tokens or _default_max_tokens,
         "temperature": _resolve_temperature(request.temperature),
         "top_p": _resolve_top_p(request.top_p),
     }
+    if rep_penalty is not None:
+        chat_kwargs["repetition_penalty"] = rep_penalty
 
     # Add multimodal content
     if has_media:
@@ -1862,15 +1887,18 @@ async def stream_completion(
     engine: BaseEngine,
     prompt: str,
     request: CompletionRequest,
+    repetition_penalty: float | None = None,
 ) -> AsyncIterator[str]:
     """Stream completion response."""
-    async for output in engine.stream_generate(
-        prompt=prompt,
-        max_tokens=request.max_tokens or _default_max_tokens,
-        temperature=_resolve_temperature(request.temperature),
-        top_p=_resolve_top_p(request.top_p),
-        stop=request.stop,
-    ):
+    gen_kwargs = {
+        "max_tokens": request.max_tokens or _default_max_tokens,
+        "temperature": _resolve_temperature(request.temperature),
+        "top_p": _resolve_top_p(request.top_p),
+        "stop": request.stop,
+    }
+    if repetition_penalty is not None:
+        gen_kwargs["repetition_penalty"] = repetition_penalty
+    async for output in engine.stream_generate(prompt=prompt, **gen_kwargs):
         data = {
             "id": f"cmpl-{uuid.uuid4().hex[:8]}",
             "object": "text_completion",

vllm_mlx/utils/tokenizer.py

@@ -59,7 +59,14 @@ def load_model_with_fallback(model_name: str, tokenizer_config: dict = None):
             return _load_with_tokenizer_fallback(model_name)
         # Fallback for models with extra weights (e.g., MTP layers)
         elif "parameters not in model" in str(e):
-            logger.warning(f"Extra parameters found (e.g., MTP weights), retrying with strict=False: {e}")
+            logger.warning(f"Extra parameters found (e.g., MTP weights), retrying with strict=False")
+            # Clear traceback references to free memory from the failed first load.
+            # Without this, large models (200GB+) cause OOM during retry because
+            # the traceback holds references to the first load's weight tensors.
+            e.__traceback__ = None
+            del e
+            import gc
+            gc.collect()
             from mlx_lm.utils import _download, load_model, load_tokenizer
             model_path = _download(model_name)
             model, config = load_model(model_path, strict=False)