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 <[email protected]>
This commit is contained in:
Mustafa Alp Yılmaz
2026-02-20 18:00:59 +00:00
committed by GitHub
co-authored by rltakashige
parent 6b54a27019
commit 6b5a705959
@@ -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: