Fix: Correctly filter special tokens in benchmark_prefix_caching #28615
Conversation
|
👋 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 You ask your reviewers to trigger select CI tests on top of 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 If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
![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 bug in the sample_tokens function within benchmarks/benchmark_prefix_caching.py. The original implementation failed to filter special tokens due to an incorrect comparison between token strings and a set of special token IDs. The fix correctly compares the token ID with the set of special IDs, resolving the bug. Furthermore, the code is made more efficient by using vocab.values() instead of vocab.items(), as the dictionary keys were unused. The change is correct, concise, and I have no further suggestions for improvement.
Signed-off-by: Di Wu <dw2761@nyu.edu>
dw2761
force-pushed
the
fix-prefix-caching-sampler
branch
from
7af6575 to
93db02b
Compare
github-actions
bot
added
the
ready
…m-project#28615) Signed-off-by: Di Wu <dw2761@nyu.edu> Signed-off-by: George D. Torres <gdavtor@gmail.com>
…m-project#28615) Signed-off-by: Di Wu <dw2761@nyu.edu> Signed-off-by: Bram Wasti <bwasti@meta.com>
2 participants
Purpose
Fixes a bug in
benchmarks/benchmark_prefix_caching.pywhere thesample_tokensfunction failed to filter special tokens.The original code was incorrectly comparing token strings (
k) with a set of special token IDs (all_special_ids). This PR fixes the logic to correctly compare the token ID (v) and simplifies the list comprehension by usingvocab.values()instead ofvocab.items(), as the token string key is not used.Test Plan
This is a self-contained logic fix within a benchmark script. The test plan is to ensure the script still runs successfully after the change.
I ran the benchmark script locally to confirm that it executes without any new errors.
Test Result
The script runs successfully. The
sample_tokensfunction now correctly filters all special token IDs from the sampling pool, as the original code intended.