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) <noreply@anthropic.com>
This commit is contained in:
dmcc73
2026-03-30 23:41:47 +01:00
parent e1cf376e45
commit b47a287f3e
@@ -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,