Skip to content

[Bugfix][V1] Fix allowed_token_ids for v1 Sampler #14169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 5, 2025
Merged

[Bugfix][V1] Fix allowed_token_ids for v1 Sampler #14169

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a5457a6
Fix allowed_token_ids for v1 Sampler
houseroad Mar 4, 2025
56cf0cc
address comments
houseroad Mar 4, 2025
830f784
address comments to make sure we set the default value for mask corre…
houseroad Mar 4, 2025
661ecf3
address comments to make sure we set the default value for mask corre…
houseroad Mar 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions vllm/v1/engine/processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,13 @@ def _validate_allowed_token_ids(
return
if params.allowed_token_ids is None:
return
if not all(0 <= tid < self.model_config.vocab_size
if params.allowed_token_ids is not None and len(
params.allowed_token_ids) == 0:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can be just if not params.allowed_token_ids since if params.allowed_token_ids is None it would have already returned above.

raise ValueError("allowed_token_ids is not None and empty!")
if not all(0 <= tid < self.model_config.get_vocab_size()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if params.allowed_token_ids is not None and len(
params.allowed_token_ids) == 0:
raise ValueError("allowed_token_ids is not None and empty!")
if not all(0 <= tid < self.model_config.get_vocab_size()
if not params.allowed_token_ids:
raise ValueError("allowed_token_ids cannot be empty")
vocab_size = self.model_config.get_vocab_size()
if not all(0 <= tid < vocab_size

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make sense, addressed the comments

for tid in params.allowed_token_ids):
raise ValueError(
"allowed_token_ids contains out-of-vocab token id")
"allowed_token_ids contains out-of-vocab token id!")

def process_inputs(
self,
Expand Down
14 changes: 7 additions & 7 deletions vllm/v1/worker/gpu_input_batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,17 +300,17 @@ def add_request(
self.has_allowed_token_ids.add(req_id)
if self.allowed_token_ids_mask_cpu_tensor is None:
# Lazy allocation for this tensor, which can be large.
self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,
self.vocab_size,
dtype=torch.bool,
device=self.device)
self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
self.allowed_token_ids_mask = torch.ones(self.max_num_reqs,
self.vocab_size,
dtype=torch.bool,
device=self.device)
self.allowed_token_ids_mask_cpu_tensor = torch.ones(
self.max_num_reqs,
self.vocab_size,
dtype=torch.bool,
device="cpu")
self.allowed_token_ids_mask_cpu_tensor[req_index][
sampling_params.allowed_token_ids] = True
sampling_params.allowed_token_ids] = False
Copy link
Member

@ywang96 ywang96 Mar 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the updated logic here is a bit counter-intuitive to readers from a glance - can we add a NOTE here that all token ids with True mask will be filled with float("-inf") during sampling?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, added some comments


# Add request lora ID
if request.lora_request:
Expand Down Expand Up @@ -359,7 +359,7 @@ def remove_request(self, req_id: str) -> Optional[int]:
self.logit_bias[req_index] = None
self.has_allowed_token_ids.discard(req_id)
if self.allowed_token_ids_mask_cpu_tensor is not None:
self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(True)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@houseroad this should also be reverted right?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

True, added more comments to help understand

return req_index

def swap_states(self, i1: int, i2: int) -> None:
Expand Down