[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

+@pytest.mark.parametrize("device", CUDA_DEVICES)
+@pytest.mark.parametrize("batch_size", [1, 2, 32])
+@pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])
+def test_sampler_allowed_token_ids(device: str, batch_size: int,


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

                sampling_metadata.presence_penalties,
                sampling_metadata.frequency_penalties,
                sampling_metadata.repetition_penalties,
-                sampling_metadata.output_token_ids)
+                sampling_metadata.output_token_ids,
+            )


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 <lufang@fb.com> Complete the main logic for allowed_token_ids in v1 sampler Signed-off-by: Lu Fang <lufang@fb.com> lint the code Signed-off-by: Lu Fang <lufang@fb.com> Add boundary checks for allowed_token_ids Signed-off-by: Lu Fang <lufang@fb.com> address the comments Signed-off-by: Lu Fang <lufang@fb.com>

Signed-off-by: Lu Fang <lufang@fb.com>

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, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat 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 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 <lufang@fb.com>

move the check in the Processor

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.

Learn about vigilant mode

Loading
Loading status checks…

60b4393

Signed-off-by: Lu Fang <lufang@fb.com>

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

fix more tests

Loading
Loading status checks…

beaba71

Signed-off-by: Lu Fang <lufang@fb.com>

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 <lufang@fb.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[v1] Support allowed_token_ids in v1 Sampler #13210

[v1] Support allowed_token_ids in v1 Sampler #13210

houseroad commented Feb 13, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Feb 13, 2025

mergify bot commented Feb 16, 2025

WoosukKwon commented Feb 17, 2025

njhill left a comment

njhill Feb 17, 2025

mergify bot commented Feb 19, 2025

njhill left a comment

njhill Feb 19, 2025

houseroad Feb 19, 2025

houseroad Feb 20, 2025

njhill left a comment

njhill left a comment

[v1] Support allowed_token_ids in v1 Sampler #13210

[v1] Support allowed_token_ids in v1 Sampler #13210

Conversation

houseroad commented Feb 13, 2025 • edited by github-actions bot Loading

github-actions bot commented Feb 13, 2025

mergify bot commented Feb 16, 2025

WoosukKwon commented Feb 17, 2025

njhill left a comment

Choose a reason for hiding this comment

njhill Feb 17, 2025

Choose a reason for hiding this comment

mergify bot commented Feb 19, 2025

njhill left a comment

Choose a reason for hiding this comment

njhill Feb 19, 2025

Choose a reason for hiding this comment

houseroad Feb 19, 2025

Choose a reason for hiding this comment

houseroad Feb 20, 2025

Choose a reason for hiding this comment

njhill left a comment

Choose a reason for hiding this comment

njhill left a comment

Choose a reason for hiding this comment

houseroad commented Feb 13, 2025 •

edited by github-actions bot

Loading