Skip to content

[BugFix] reschedule_preempt_task append waiting & PREEMPTED blocksize #5506

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jiang-Jia-Jun merged 6 commits into from
Dec 12, 2025
Merged

[BugFix] reschedule_preempt_task append waiting & PREEMPTED blocksize #5506

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ad1e54c
bugfix reschedule_preempt_task append waiting & PREEMPTED blocksize
ST-XX Dec 11, 2025
fe544da
bugfix reschedule_preempt_task append waiting & PREEMPTED blocksize
ST-XX Dec 11, 2025
0c81631
注释
ST-XX Dec 11, 2025
a7bf030
[bugfix] PREEMPTED task blocksize
ST-XX Dec 11, 2025
f5053b9
Apply suggestion from @Copilot
ST-XX Dec 12, 2025
e640adf
Merge branch 'develop' into feature/preempt_task
ST-XX Dec 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions fastdeploy/engine/sched/resource_manager_v1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,7 @@ def _trigger_preempt(self, request, num_new_blocks, preempted_reqs, scheduled_re
self.running.insert(0, preempted_req)
continue
preempted_req.status = RequestStatus.PREEMPTED
preempted_req.last_preempted_blocksize = len(preempted_req.block_tables)
preempted_req.num_computed_tokens = 0
if self.config.scheduler_config.splitwise_role == "decode":
self.tasks_list[preempted_req.idx] = None
Expand Down Expand Up @@ -735,6 +736,14 @@ def _allocate_decode_and_extend():
break
num_new_tokens = self._get_num_new_tokens(request, token_budget)
num_new_block = self.get_new_block_nums(request, num_new_tokens)
# If num_new_block is less than the last preempted block size, use the last preempted block size instead.
# For normal requests, when allocating blocks, we reserve two extra blocks for decoding.
# In the request rescheduling scenario, we currently only consider the number of tokens already generated,
# which might lead to allocating fewer blocks than the previous allocation, causing repeated rescheduling.
# This adjustment ensures we at least allocate as many blocks as before to avoid this issue.
last_preempted_blocksize = getattr(request, "last_preempted_blocksize", 0)
if num_new_block < last_preempted_blocksize:
num_new_block = last_preempted_blocksize
# Allocate blocks to prefill
if self.cache_manager.can_allocate_gpu_blocks(num_new_block):
if not request.get("skip_allocate", False):
Expand Down
Loading