skip nan logprobs on tokens (#1622)

c4e874e97db6e11737e076675dd1e112e10ea7f5 · 2026-02-25 17:44:15 +0000 · Evan Quiney

sometimes we generate NaN logprobs for tokens, this causes pydantic
validation errors on the receiving end. in this case we should just not
send the logprob items

Files touched

M src/exo/worker/engines/mlx/generator/generate.py

Diff

commit c4e874e97db6e11737e076675dd1e112e10ea7f5
Author: Evan Quiney <evanev7@gmail.com>
Date:   Wed Feb 25 17:44:15 2026 +0000

    skip nan logprobs on tokens (#1622)
    
    sometimes we generate NaN logprobs for tokens, this causes pydantic
    validation errors on the receiving end. in this case we should just not
    send the logprob items
---
 src/exo/worker/engines/mlx/generator/generate.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 6e926205..e7307b8d 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -1,3 +1,4 @@
+import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
@@ -248,6 +249,9 @@ def extract_top_logprobs(
     for i in range(top_logprobs):
         token_id = int(top_indices[i].item())
         token_logprob = float(top_values[i].item())
+        if math.isnan(token_logprob):
+            continue
+
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
         # Get byte representation

← e23c3a30 Address Mac Mini pipeline GPU timeouts (#1620) · back to Exo · report macmon failures more aggressively (#1618) eab3e0b4 →