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