From b47a287f3e0099bcee1b5265c1502c1fc4d3000f Mon Sep 17 00:00:00 2001 From: dmcc73 Date: Mon, 30 Mar 2026 23:41:47 +0100 Subject: [PATCH] Add EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction For profiling: extract_top_logprobs() does 11 .item() calls + argpartition on 248K vocab per token. Testing if this is the source of speculative overhead vs mlx_bench. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/exo/worker/engines/mlx/generator/batch_generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 256d5837..64216ffd 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -470,7 +470,7 @@ class ExoBatchGenerator: logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None - if task_params.logprobs: + if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1": logprob, top_logprobs = extract_top_logprobs( logprobs=response.logprobs, tokenizer=self.tokenizer,