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

@@ -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 }

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()

src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py

@@ -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)

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'" },