@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8d6cd32b7b

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

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

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

Lu Fang committed Mar 5, 2025 at 00:49 UTC 8d6cd32b7b4b80207ba210d9e805d07f00bfde07

2 files changed +12 -4

vllm/v1/engine/processor.py

+5 -3

@@ -92,10 +92,12 @@ class Processor:

92 return

93 if params.allowed_token_ids is None:

94 return

95 - if not all(0 <= tid < self.model_config.vocab_size

96 - for tid in params.allowed_token_ids):

95 + if not params.allowed_token_ids:

96 + raise ValueError("allowed_token_ids is not None and empty!")

97 + vocab_size = self.model_config.get_vocab_size()

98 + if not all(0 <= tid < vocab_size for tid in params.allowed_token_ids):

99 raise ValueError(

98 - "allowed_token_ids contains out-of-vocab token id")

100 + "allowed_token_ids contains out-of-vocab token id!")

101

102 def process_inputs(

103 self,

vllm/v1/worker/gpu_input_batch.py

+7 -1

@@ -199,6 +199,8 @@ class InputBatch:

199 self.logit_bias: list[Optional[dict[int,

200 float]]] = [None] * max_num_reqs

201 self.has_allowed_token_ids: set[str] = set()

202 + # NOTE(lufang): In the mask tensor, if the corresponding token allowed,

203 + # the value is False. Since we use masked_fill_ to set -inf.

204 self.allowed_token_ids_mask: Optional[torch.Tensor] = None

205 self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None

206

@@ -300,6 +302,7 @@ class InputBatch:

302 self.has_allowed_token_ids.add(req_id)

303 if self.allowed_token_ids_mask_cpu_tensor is None:

304 # Lazy allocation for this tensor, which can be large.

305 + # False means we don't fill with -inf.

306 self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,

307 self.vocab_size,

308 dtype=torch.bool,

@@ -309,8 +312,10 @@ class InputBatch:

312 self.vocab_size,

313 dtype=torch.bool,

314 device="cpu")

315 + self.allowed_token_ids_mask_cpu_tensor[req_index] = True

316 + # False means we don't fill with -inf.

317 self.allowed_token_ids_mask_cpu_tensor[req_index][

313 - sampling_params.allowed_token_ids] = True

318 + sampling_params.allowed_token_ids] = False

319

320 # Add request lora ID

321 if request.lora_request:

@@ -359,6 +364,7 @@ class InputBatch:

364 self.logit_bias[req_index] = None

365 self.has_allowed_token_ids.discard(req_id)

366 if self.allowed_token_ids_mask_cpu_tensor is not None:

367 + # False means we don't fill with -inf.

368 self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)

369 return req_index

370