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) <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.6
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,
|
||||
|
||||
Reference in New Issue
Block a user