fix: immediate cancel check after prefill completes (#1575)

6b5a705959708c221cd41444cbfbf9c8360d5da4 · 2026-02-20 21:00:59 +0300 · Mustafa Alp Yılmaz

## 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 <rl.takashige@gmail.com>

Files touched

M src/exo/worker/runner/llm_inference/runner.py

Diff

commit 6b5a705959708c221cd41444cbfbf9c8360d5da4
Author: Mustafa Alp Yılmaz <96022931+mustafalpyilmaz@users.noreply.github.com>
Date:   Fri Feb 20 21:00:59 2026 +0300

    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 <rl.takashige@gmail.com>
---
 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:

← 6b54a270 fix: add downloaded_bytes to DownloadPending event (#1564) · back to Exo · worker: add EXO_MODELS_PATH for pre-downloaded model directo 42da58c2 →