From 6b5a705959708c221cd41444cbfbf9c8360d5da4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?= <96022931+mustafalpyilmaz@users.noreply.github.com> Date: Fri, 20 Feb 2026 21:00:59 +0300 Subject: [PATCH] fix: immediate cancel check after prefill completes (#1575) ## Problem When a request is cancelled during prefill, the cancellation is not detected until `check_for_cancel_every` additional tokens have been generated. This is because `tokens_since_last_cancel_check` is initialized to `0`, meaning the first cancel check only happens after generating `check_for_cancel_every` tokens post-prefill. For long prefills (which are the most likely to be cancelled), this adds unnecessary latency before the cancellation is actually honoured. ## Fix Initialize `tokens_since_last_cancel_check` to `check_for_cancel_every` instead of `0`, so the very first token generated after prefill triggers an immediate cancel check. ```diff - tokens_since_last_cancel_check = 0 + tokens_since_last_cancel_check = check_for_cancel_every ``` ## Impact - Cancellations issued during prefill are detected immediately when generation begins - No change in behaviour for non-cancelled requests (the counter resets to `0` after each check as before) - 1 line changed Co-authored-by: rltakashige --- src/exo/worker/runner/llm_inference/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index 5581f6dc..c4059b9e 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -313,7 +313,7 @@ def main( mlx_generator = parse_tool_calls(mlx_generator, tool_parser) completion_tokens = 0 - tokens_since_last_cancel_check = 0 + tokens_since_last_cancel_check = check_for_cancel_every for response in mlx_generator: tokens_since_last_cancel_check += 1 if tokens_since_last_cancel_check >= check_for_cancel_every: