Skip to content

fix: validate max_new_tokens against speculative_tokens during warmup (fixes #474) - #793

Open
mithilesh1909 wants to merge 1 commit into
predibase:mainfrom
mithilesh1909:fix/issue-474-speculative-warmup-check
Open

fix: validate max_new_tokens against speculative_tokens during warmup (fixes #474)#793
mithilesh1909 wants to merge 1 commit into
predibase:mainfrom
mithilesh1909:fix/issue-474-speculative-warmup-check

Conversation

@mithilesh1909

Copy link
Copy Markdown

… (fixes #474)

What does this PR do?

When speculative decoding is enabled, generating up to the model's max positional embeddings causes a CUDA device-side assert error at runtime. This PR adds an early validation check in the router that rejects requests where max_new_tokens + speculative_tokens would exceed the available token budget (max_total_tokens - input_length), returning a clear 422 error instead of crashing with a CUDA error.

Changes:

  • Added speculative_tokens field to the Validation struct
  • Added MaxNewTokensSpeculative error variant with a descriptive message
  • Added validation check in both tokenized and non-tokenized input paths
  • Added a unit test covering the boundary conditions
  • Updated Validation::new() call in server.rs to pass speculative_tokens

Fixes #474

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

…fixes predibase#474)

Signed-off-by: mithilesh1909 <mittuayodhya@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Improve warmup checking for max new tokens when using speculative decoding

1 participant