[v1] Support allowed_token_ids in v1 Sampler #13210

houseroad · 2025-02-13T07:31:47Z

Follow the implementation in vllm/entrypoints/openai/logits_processors.py.

The idea is straightforward, adding a [batch_size x vocab_size] mask tensor, and leverage a list of bools to determine whether to do the inplace masked fill.

add test
move some verification in _get_allowed_token_ids_logits_processor to SamplingParam.

Test with

pytest tests/v1/sample/test_sampler.py
pytest tests/v1/worker/test_gpu_input_batch.py

github-actions · 2025-02-13T07:31:57Z

👋 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 fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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 ready label to the PR or enable auto-merge.

🚀

mergify · 2025-02-16T00:27:59Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @houseroad.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

WoosukKwon · 2025-02-17T09:00:36Z

cc @njhill Could you please take a look?

njhill

Thanks @houseroad!

I have a few comments, mainly to keep the tensor manipulation on cpu and apply mask to the batch rather than individual rows.

Also you might want to take a look at #13360, I'm hoping to get that into shape soon and have all the sampling params like this encapsulated as vectorized logits processors.

vllm/v1/sample/sampler.py

vllm/v1/worker/gpu_input_batch.py

vllm/v1/sample/sampler.py

njhill · 2025-02-17T18:46:01Z

tests/v1/sample/test_sampler.py

Thank you for adding this test!!

mergify · 2025-02-19T02:07:05Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @houseroad.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

njhill

Thanks @lu-wang-dl for the updates, I think it's close, just a few more things.

vllm/v1/sample/metadata.py

njhill · 2025-02-19T21:22:48Z

vllm/v1/sample/sampler.py

Same general comment, it's best not to include unrelated formatting changes, looks like there are a couple above and some in the other file.

yeah, it's done by my VSCode, sorry.

I will be more cautious next time.

okay, confirmed, actually it's done by the yapf in the pre-commit hooks. It forced me to do such change...

vllm/v1/sample/sampler.py

vllm/v1/worker/gpu_input_batch.py

njhill

Thanks @houseroad looks great now, just one remaining thing.

vllm/v1/worker/gpu_input_batch.py

njhill

Thanks again @houseroad!

Signed-off-by: Lu Fang <[email protected]> Complete the main logic for allowed_token_ids in v1 sampler Signed-off-by: Lu Fang <[email protected]> lint the code Signed-off-by: Lu Fang <[email protected]> Add boundary checks for allowed_token_ids Signed-off-by: Lu Fang <[email protected]> address the comments Signed-off-by: Lu Fang <[email protected]>

Signed-off-by: Lu Fang <[email protected]>

Signed-off-by: Lu Fang <[email protected]> Signed-off-by: Louis Ulmer <[email protected]>

Signed-off-by: Lu Fang <[email protected]>

mergify bot added the v1 label Feb 13, 2025

houseroad force-pushed the v1_allowed_token_ids branch from 9b044a4 to ee1d1a0 Compare February 14, 2025 22:16

houseroad changed the title ~~Support allowed_token_ids in v1 Sampler~~ [v1] Support allowed_token_ids in v1 Sampler Feb 14, 2025

mergify bot added the needs-rebase label Feb 16, 2025

houseroad force-pushed the v1_allowed_token_ids branch from 3a5f408 to 1d7bcfc Compare February 16, 2025 04:53

mergify bot removed the needs-rebase label Feb 16, 2025

houseroad marked this pull request as ready for review February 16, 2025 05:32

houseroad requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners February 16, 2025 05:32

njhill reviewed Feb 17, 2025

View reviewed changes

njhill mentioned this pull request Feb 18, 2025

[V1][core] Support allowed_token_ids in V1 sampler #13481

Closed

mergify bot added the needs-rebase label Feb 19, 2025

houseroad force-pushed the v1_allowed_token_ids branch from 1793e7a to 77c8db6 Compare February 19, 2025 05:39

mergify bot removed the needs-rebase label Feb 19, 2025

houseroad force-pushed the v1_allowed_token_ids branch 3 times, most recently from 9c0d832 to 2bdd559 Compare February 19, 2025 06:56

WoosukKwon mentioned this pull request Feb 16, 2025

[V1][Help Wanted] Porting missing sampling parameters to V1 #13058

Closed

7 tasks

njhill reviewed Feb 19, 2025

View reviewed changes

njhill reviewed Feb 20, 2025

View reviewed changes

vllm/v1/worker/gpu_input_batch.py Outdated Show resolved Hide resolved

njhill approved these changes Feb 21, 2025

View reviewed changes

houseroad closed this Feb 21, 2025

houseroad reopened this Feb 21, 2025

ywang96 added the ready ONLY add when PR is ready to merge/full CI is needed label Feb 21, 2025

houseroad force-pushed the v1_allowed_token_ids branch from 444d719 to 7c63284 Compare February 21, 2025 05:54

houseroad added 3 commits February 21, 2025 12:04

address comments

368369c

Signed-off-by: Lu Fang <[email protected]>

move the check in the Processor

60b4393

Signed-off-by: Lu Fang <[email protected]>

houseroad force-pushed the v1_allowed_token_ids branch from 7c63284 to 60b4393 Compare February 21, 2025 20:04

fix more tests

beaba71

Signed-off-by: Lu Fang <[email protected]>

DarkLight1337 merged commit bb78fb3 into vllm-project:main Feb 22, 2025
42 checks passed

Akshat-Tripathi pushed a commit to krai/vllm that referenced this pull request Mar 3, 2025

[v1] Support allowed_token_ids in v1 Sampler (vllm-project#13210)

b8466ec

Signed-off-by: Lu Fang <[email protected]>

lulmer pushed a commit to lulmer/vllm that referenced this pull request Apr 7, 2025

[v1] Support allowed_token_ids in v1 Sampler (vllm-project#13210)

fc067f7

Signed-off-by: Lu Fang <[email protected]> Signed-off-by: Louis Ulmer <[email protected]>

ckhordiasma mentioned this pull request Apr 17, 2025

[do not merge] pr test for nm changes into 2.20 red-hat-data-services/vllm#107

Closed

shreyankg pushed a commit to shreyankg/vllm that referenced this pull request May 3, 2025

[v1] Support allowed_token_ids in v1 Sampler (vllm-project#13210)

588f558

Signed-off-by: Lu Fang <[email protected]>

Uh oh!

[v1] Support allowed_token_ids in v1 Sampler #13210

[v1] Support allowed_token_ids in v1 Sampler #13210

Uh oh!

Conversation

houseroad commented Feb 13, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Feb 13, 2025

Uh oh!

mergify bot commented Feb 16, 2025

Uh oh!

WoosukKwon commented Feb 17, 2025

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

njhill Feb 17, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Feb 19, 2025

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

njhill Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

houseroad Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

houseroad Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

houseroad commented Feb 13, 2025 •

edited by github-actions bot

Loading