move the check in the Processor

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

Author: houseroad

vllm/v1/engine/processor.py

@@ -83,6 +83,19 @@ def _validate_lora(self, lora_request: Optional[LoRARequest]) -> None:
             raise ValueError(f"Got lora_request {lora_request} but LoRA is "
                              "not enabled!")
 
+    def _validate_allowed_token_ids(
+        self,
+        params: Union[SamplingParams, PoolingParams],
+    ) -> None:
+        if not isinstance(params, SamplingParams):
+            return
+        if params.allowed_token_ids is None:
+            return
+        if not all(0 <= tid < self.model_config.vocab_size
+                   for tid in params.allowed_token_ids):
+            raise ValueError(
+                "allowed_token_ids contains out-of-vocab token id")
+
     def process_inputs(
         self,
         request_id: str,
@@ -100,6 +113,7 @@ def process_inputs(
 
         self._validate_logprobs(params)
         self._validate_lora(lora_request)
+        self._validate_allowed_token_ids(params)
 
         if arrival_time is None:
             arrival_time = time.time()

vllm/v1/sample/sampler.py

@@ -58,8 +58,8 @@ def forward(
 
         # Gather the logprobs of the topk and sampled token (if requested).
         # Get logprobs and rank tensors (if requested)
-        logprobs_tensors = (None if num_logprobs is None else \
-            self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=sampled))
+        logprobs_tensors = None if num_logprobs is None else \
+            self.gather_logprobs(raw_logprobs, num_logprobs, token_ids=sampled)
 
         # Use int32 to reduce the tensor size.
         sampled = sampled.to(torch.int32)
@@ -183,7 +183,6 @@ def apply_penalties(
             apply_min_token_penalties(logits,
                                       sampling_metadata.output_token_ids,
                                       sampling_metadata.min_tokens)
-
         if not sampling_metadata.no_penalties:
             assert sampling_metadata.prompt_token_ids is not None
             logits = apply_all_penalties(

vllm/v1/worker/gpu_input_batch.py

@@ -291,12 +291,6 @@ def add_request(
             self.logit_bias[req_index] = sampling_params.logit_bias
 
         if sampling_params.allowed_token_ids:
-            # NOTE(houseroad): put the check here since no vocab_size info
-            # available in vllm/sampling_params.py
-            if not all(0 <= tid < self.vocab_size
-                       for tid in sampling_params.allowed_token_ids):
-                raise ValueError(
-                    "allowed_token_ids contains out-of-vocab token id")
             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.