[object Object]

← back to Exo

Prevent Qwen3.5 looping by using mlx lm fork (#1784)

b6240a97e896fe09e8b402601d2ef2683e4d0de4 · 2026-03-24 17:16:41 +0000 · rltakashige

## Motivation

Move back to an MLX LM to improve the Qwen 3.5 experience. 

## Test Plan

### Manual Testing
Seems to loop less from testing, no speed regressions.

Files touched

Diff

commit b6240a97e896fe09e8b402601d2ef2683e4d0de4
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Mar 24 17:16:41 2026 +0000

    Prevent Qwen3.5 looping by using mlx lm fork (#1784)
    
    ## Motivation
    
    Move back to an MLX LM to improve the Qwen 3.5 experience.
    
    ## Test Plan
    
    ### Manual Testing
    Seems to loop less from testing, no speed regressions.
---
 pyproject.toml                                     |  2 +-
 src/exo/worker/engines/mlx/patches/__init__.py     |  5 +---
 .../mlx/patches/high_precision_gdn_softplus.py     | 27 ----------------------
 uv.lock                                            |  4 ++--
 4 files changed, 4 insertions(+), 34 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index bab185d8..de5238c9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "fix/float32-logprobs" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
diff --git a/src/exo/worker/engines/mlx/patches/__init__.py b/src/exo/worker/engines/mlx/patches/__init__.py
index 0a14b262..b05a045e 100644
--- a/src/exo/worker/engines/mlx/patches/__init__.py
+++ b/src/exo/worker/engines/mlx/patches/__init__.py
@@ -1,6 +1,3 @@
-from exo.worker.engines.mlx.patches.high_precision_gdn_softplus import (
-    patch_gdn_softplus,
-)
 from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
 from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
 
@@ -13,5 +10,5 @@ def apply_mlx_patches() -> None:
         return
     _applied = True
     patch_yarn_rope()
-    patch_gdn_softplus()
+    # patch_gdn_softplus()
     apply_batch_gen_patch()
diff --git a/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py b/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py
deleted file mode 100644
index a1961e11..00000000
--- a/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import sys
-
-import mlx.core as mx
-from mlx_lm.models.gated_delta import compute_g
-
-
-def _compute_g_f32(a_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
-    return mx.exp(
-        -mx.exp(a_log.astype(mx.float32))
-        * mx.where(
-            (a + dt_bias).astype(mx.float32) > 20,
-            (a + dt_bias).astype(mx.float32),
-            mx.log1p(mx.exp((a + dt_bias).astype(mx.float32))),
-        )
-    ).astype(a.dtype)
-
-
-def patch_gdn_softplus() -> None:
-    from mlx_lm.models import gated_delta
-
-    gated_delta.compute_g = _compute_g_f32
-
-    for mod in list(sys.modules.values()):
-        if mod is gated_delta:
-            continue
-        if getattr(mod, "compute_g", None) is compute_g:
-            object.__setattr__(mod, "compute_g", _compute_g_f32)
diff --git a/uv.lock b/uv.lock
index 82411d80..9579f581 100644
--- a/uv.lock
+++ b/uv.lock
@@ -524,7 +524,7 @@ requires-dist = [
     { name = "mflux", specifier = "==0.16.9" },
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?branch=main" },
+    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=fix%2Ffloat32-logprobs" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "psutil", specifier = ">=7.0.0" },
@@ -1445,7 +1445,7 @@ wheels = [
 [[package]]
 name = "mlx-lm"
 version = "0.31.2"
-source = { git = "https://github.com/ml-explore/mlx-lm?branch=main#ed7884cb80968e0e77fce6cde5d1597952bbd524" }
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=fix%2Ffloat32-logprobs#8e94256220f954949133e036980951681e353945" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },

← 6cdfbb7e Add HF_ENDPOINT in the app settings (#1783)  ·  back to Exo  ·  send error finish reason on failing to parse a tool call (#1 509533d4 →