[BugFix] Fix spec decoding max_tokens scheduling perf issue #29542
Conversation
This affected performance with async scheduling since requests could be unscheduled before they were finished and have to be scheduled again for the last 1-2 tokens. Signed-off-by: Nick Hill <nhill@redhat.com>
njhill
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
robertgshaw2-redhat and
ywang96
as code owners
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses a performance issue related to speculative decoding with asynchronous scheduling. The change correctly adjusts the calculation of max_total_tokens to account for speculative placeholder tokens. This prevents requests from being prematurely unscheduled when they are near their max_tokens limit, which previously required them to be rescheduled to generate the final tokens. The logic is sound, ensuring that enough tokens can be scheduled for speculative verification while still respecting the model's maximum length. The implementation is clear and well-commented, and it should effectively resolve the described performance bottleneck.
Signed-off-by: Nick Hill <nhill@redhat.com>
| start = start_req_idx | ||
| end = end_req_idx | ||
| sliced_cu_num_generated_tokens = None | ||
| def slice_request(self, req_idx: int, num_positions: int): |
There was a problem hiding this comment.
Changing how logprobs are sliced to ensure we don't include extra positions beyond the number of output tokens.
2 participants
This affected performance with async scheduling since requests could be unscheduled before they were finished and have to be scheduled again for the last 1-2 tokens.
Thanks @benchislett for finding.