[v1] Support allowed_token_ids in v1 Sampler #13210
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 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 🚀 |
houseroad
force-pushed
the
v1_allowed_token_ids
branch
from
9b044a4 to
ee1d1a0
Compare
houseroad
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
houseroad
force-pushed
the
v1_allowed_token_ids
branch
from
3a5f408 to
1d7bcfc
Compare
houseroad
requested review from
WoosukKwon,
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
|
cc @njhill Could you please take a look? |
There was a problem hiding this comment.
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
Outdated
| logprobs_tensors = (None | ||
| if num_logprobs is None else self.gather_logprobs( | ||
| raw_logprobs, num_logprobs, token_ids=sampled)) |
There was a problem hiding this comment.
Please revert unrelated formatting changes
vllm/v1/sample/sampler.py
Outdated
| # Apply temperature. | ||
| logits = self.apply_temperature(logits, sampling_metadata.temperature) |
There was a problem hiding this comment.
This shouldn't be here (guessing it was a merge mistake)
vllm/v1/worker/gpu_input_batch.py
Outdated
| self.allowed_token_ids_mask: torch.Tensor = torch.zeros( | ||
| max_num_reqs, self.vocab_size, dtype=torch.bool, device=device) |
There was a problem hiding this comment.
This could actually get quite large (e.g. 128 MB), we may want to consider a different (pre)allocation strategy.
Also, consistent with the other tensors, we should keep this in pinned CPU memory and only copy to GPU in an async manner when the batch makeup changes (and has_allowed_token_ids is True)
vllm/v1/sample/sampler.py
Outdated
| for i, allowed_token_ids in enumerate( | ||
| sampling_metadata.has_allowed_token_ids): | ||
| if allowed_token_ids: | ||
| logits[i].masked_fill_( | ||
| sampling_metadata.allowed_token_ids_mask[i], float("-inf")) |
There was a problem hiding this comment.
Can we do a single operation here rather than a loop (make has_allowed_token_ids a single bool rather than list).
| @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, |
There was a problem hiding this comment.
Thank you for adding this test!!
|
This pull request has merge conflicts that must be resolved before it can be |
houseroad
force-pushed
the
v1_allowed_token_ids
branch
from
1793e7a to
77c8db6
Compare
houseroad
force-pushed
the
v1_allowed_token_ids
branch
2 times, most recently
from
14b0933 to
9c0d832
Compare
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]>
houseroad
force-pushed
the
v1_allowed_token_ids
branch
from
9c0d832 to
2bdd559
Compare
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.
Test with