[object Object]

← back to Exo

feat(mlx): add repetition_penalty and repetition_context_size to chat completions (#1665)

eee3432738d3f125a5889358b1177c1121a20a6c · 2026-03-06 13:51:25 +0300 · Mustafa Alp Yılmaz

## Problem

Models running through exo can get stuck in repetition loops —
generating the same text over and over until hitting the token limit. It
happens more with quantized models where probability distributions can
become degenerate.

`mlx_lm` has `make_logits_processors(repetition_penalty,
repetition_context_size)` that handles this, but it is never called
anywhere in the pipeline.

## Solution

Wire `repetition_penalty` and `repetition_context_size` through the
stack:

- `ChatCompletionRequest` → accept the params from the client
- `TextGenerationTaskParams` → carry them through the pipeline
- `chat_completions.py` adapter → map to internal params
- `generate.py` → call `make_logits_processors()` and merge into the
`logits_processors` list

When `repetition_penalty` is `None` (default), `make_logits_processors`
returns `[]` so existing behavior is unchanged.

## Usage

```json
{
  "model": "mlx-community/...",
  "messages": [...],
  "repetition_penalty": 1.15,
  "repetition_context_size": 64
}
```

Files touched

Diff

commit eee3432738d3f125a5889358b1177c1121a20a6c
Author: Mustafa Alp Yılmaz <96022931+mustafalpyilmaz@users.noreply.github.com>
Date:   Fri Mar 6 13:51:25 2026 +0300

    feat(mlx): add repetition_penalty and repetition_context_size to chat completions (#1665)
    
    ## Problem
    
    Models running through exo can get stuck in repetition loops —
    generating the same text over and over until hitting the token limit. It
    happens more with quantized models where probability distributions can
    become degenerate.
    
    `mlx_lm` has `make_logits_processors(repetition_penalty,
    repetition_context_size)` that handles this, but it is never called
    anywhere in the pipeline.
    
    ## Solution
    
    Wire `repetition_penalty` and `repetition_context_size` through the
    stack:
    
    - `ChatCompletionRequest` → accept the params from the client
    - `TextGenerationTaskParams` → carry them through the pipeline
    - `chat_completions.py` adapter → map to internal params
    - `generate.py` → call `make_logits_processors()` and merge into the
    `logits_processors` list
    
    When `repetition_penalty` is `None` (default), `make_logits_processors`
    returns `[]` so existing behavior is unchanged.
    
    ## Usage
    
    ```json
    {
      "model": "mlx-community/...",
      "messages": [...],
      "repetition_penalty": 1.15,
      "repetition_context_size": 64
    }
    ```
---
 .mlx_typings/mlx_lm/sample_utils.pyi             |  2 +-
 src/exo/master/adapters/chat_completions.py      |  2 ++
 src/exo/shared/types/api.py                      |  2 ++
 src/exo/shared/types/text_generation.py          |  2 ++
 src/exo/worker/engines/mlx/generator/generate.py | 11 ++++++++---
 5 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/.mlx_typings/mlx_lm/sample_utils.pyi b/.mlx_typings/mlx_lm/sample_utils.pyi
index bc6955a7..c001e3ec 100644
--- a/.mlx_typings/mlx_lm/sample_utils.pyi
+++ b/.mlx_typings/mlx_lm/sample_utils.pyi
@@ -48,7 +48,7 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
-):  # -> list[Any]:
+) -> list[Callable[[mx.array, mx.array], mx.array]]:
     """
     Make logits processors for use with ``generate_step``.
 
diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py
index dd76e6b8..af13fdb0 100644
--- a/src/exo/master/adapters/chat_completions.py
+++ b/src/exo/master/adapters/chat_completions.py
@@ -104,6 +104,8 @@ def chat_request_to_text_generation(
         else None,
         logprobs=request.logprobs or False,
         top_logprobs=request.top_logprobs,
+        repetition_penalty=request.repetition_penalty,
+        repetition_context_size=request.repetition_context_size,
     )
 
 
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index e6420a9b..2f13d686 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -201,6 +201,8 @@ class ChatCompletionRequest(BaseModel):
     tools: list[dict[str, Any]] | None = None
     reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
+    repetition_penalty: float | None = None
+    repetition_context_size: int | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index 51510b5c..224f51ee 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -67,3 +67,5 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     enable_thinking: bool | None = None
     logprobs: bool = False
     top_logprobs: int | None = None
+    repetition_penalty: float | None = None
+    repetition_context_size: int | None = None
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index d33d35b3..e5e89c7b 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -10,7 +10,7 @@ from mlx_lm.generate import (
     stream_generate,
 )
 from mlx_lm.models.cache import ArraysCache, RotatingKVCache
-from mlx_lm.sample_utils import make_sampler
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.api import (
@@ -470,11 +470,16 @@ def mlx_generate(
                 f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
             )
 
-    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
+        make_logits_processors(
+            repetition_penalty=task.repetition_penalty,
+            repetition_context_size=task.repetition_context_size,
+        )
+    )
     if is_bench:
         # Only sample length eos tokens
         eos_ids = eos_ids_from_tokenizer(tokenizer)
-        logits_processors = [ban_token_ids(eos_ids)]
+        logits_processors = [ban_token_ids(eos_ids)] + logits_processors
 
     sampler = make_sampler(
         temp=task.temperature if task.temperature is not None else 0.7,

← e8c3a873 feat(dashboard): improve model picker feedback and HuggingFa  ·  back to Exo  ·  docs: Update documentation for v1.0.68 release (#1667) 7a36d396 →