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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user