[V1][Spec Decode] Optimize N-gram matching with Numba #13365
Conversation
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
WoosukKwon
requested review from
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
|
cc @LiuXiaoxuanPKU This PR is ready. Could you please take a look? |
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
| @@ -1,6 +1,7 @@ | |||
| psutil | |||
| sentencepiece # Required for LLaMA tokenizer. | |||
| numpy < 2.0.0 | |||
| numba == 0.60.0 # v0.61 doesn't support Python 3.9. Required for N-gram speculative decoding. | |||
There was a problem hiding this comment.
Shouldn't this be in requirements-cuda.txt rather than common?
There was a problem hiding this comment.
Oh I'm ok with either; I just thought it would be eventually used by others as well. Please feel free to submit a PR to move it to requirements-cuda.txt and probably requirements-rocm.txt.
|
Very excited about this! |
|
@michaelfeil Thanks! Happy to see you again :) |
…3365) Signed-off-by: Woosuk Kwon <[email protected]>
…3365) Signed-off-by: Woosuk Kwon <[email protected]>
…3365) Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Louis Ulmer <[email protected]>
…3365) Signed-off-by: Woosuk Kwon <[email protected]>
This PR optimizes the N-gram matching algorithm by JIT compiling it with Numba.
I've observed 20-30x speedup with large batch sizes: For ShareGPT benchmark with 5K requests, the cumulative overhead reduces from 54.3 sec to 1.9 sec, which is ~2.5% of the entire running time.