address comments

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

Author: houseroad

tests/v1/sample/test_sampler.py

@@ -62,9 +62,8 @@ def _create_allowed_token_ids(
     vocab_size: int,
     num_allowed_token_ids: int,
     device: torch.device,
-) -> Tuple[bool, Optional[torch.Tensor]]:
+) -> Optional[torch.Tensor]:
     mask: Optional[torch.Tensor] = None
-    no_allowed_token_ids = True
     for i in range(batch_size):
         if i % 2 == 1:
             continue
@@ -75,8 +74,7 @@ def _create_allowed_token_ids(
         start = min(i, vocab_size - 1)
         end = min(i + num_allowed_token_ids, vocab_size - 1)
         mask[i, start:end] = True
-        no_allowed_token_ids = False
-    return (no_allowed_token_ids, mask)
+    return mask
 
 
 def _create_default_sampling_metadata(
@@ -114,7 +112,6 @@ def _create_default_sampling_metadata(
         no_penalties=True,
         min_tokens={},
         logit_bias=[None] * batch_size,
-        no_allowed_token_ids=True,
         allowed_token_ids_mask=None,
     )
     return fake_sampling_metadata
@@ -448,18 +445,16 @@ def test_sampler_allowed_token_ids(device: str, batch_size: int,
     fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)
     sampling_metadata = _create_default_sampling_metadata(
         NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))
-    no_allowed_token_ids, mask = _create_allowed_token_ids(
+    mask = _create_allowed_token_ids(
         batch_size=batch_size,
         vocab_size=VOCAB_SIZE,
         num_allowed_token_ids=num_allowed_token_ids,
         device=device,
     )
-    sampling_metadata.no_allowed_token_ids = no_allowed_token_ids
     sampling_metadata.allowed_token_ids_mask = mask
     sampler = Sampler()
     logits = sampler.apply_allowed_token_ids(fake_logits, sampling_metadata)
     logits = logits.cpu()
-    assert not sampling_metadata.no_allowed_token_ids
     for batch_idx in range(batch_size):
         logits_for_req = logits[batch_idx]
         if batch_idx % 2 == 1:

tests/v1/worker/test_gpu_input_batch.py

@@ -66,7 +66,6 @@ def _construct_expected_sampling_metadata(
     temperature = [0.0 for _ in range(num_reqs)]
     min_tokens = {}
     logit_bias = [None] * num_reqs
-    has_allowed_token_ids = [False] * num_reqs
     allowed_token_ids_mask = torch.zeros(num_reqs,
                                          VOCAB_SIZE,
                                          dtype=torch.bool,
@@ -92,7 +91,6 @@ def _construct_expected_sampling_metadata(
             req.sampling_params.all_stop_token_ids)
         logit_bias[index_in_input_batch] = req.sampling_params.logit_bias
         if req.sampling_params.allowed_token_ids:
-            has_allowed_token_ids[index_in_input_batch] = True
             allowed_token_ids_mask[index_in_input_batch][
                 req.sampling_params.allowed_token_ids] = True
 
@@ -131,7 +129,6 @@ def _construct_expected_sampling_metadata(
                       and all(x == 0 for x in frequency_penalties)
                       and all(x == 1 for x in repetition_penalties)),
         logit_bias=logit_bias,
-        no_allowed_token_ids=not any(has_allowed_token_ids),
         allowed_token_ids_mask=allowed_token_ids_mask,
     )
 
@@ -254,9 +251,7 @@ def same(t1: Optional[torch.Tensor], t2: Optional[torch.Tensor]) -> bool:
     assert expected_sampling_metadata.no_penalties == \
            sampling_metadata.no_penalties
     assert expected_sampling_metadata.logit_bias == sampling_metadata.logit_bias
-    assert (expected_sampling_metadata.no_allowed_token_ids ==
-            sampling_metadata.no_allowed_token_ids)
-    if not sampling_metadata.no_allowed_token_ids:
+    if sampling_metadata.allowed_token_ids_mask:
         assert torch.allclose(
             expected_sampling_metadata.allowed_token_ids_mask,
             sampling_metadata.allowed_token_ids_mask)

vllm/v1/sample/metadata.py

@@ -38,10 +38,6 @@ class SamplingMetadata:
 
     logit_bias: List[Optional[Dict[int, float]]]
 
-    # These two parameters are for allowed_token_ids.
-    # `no_allowed_token_ids`` is a bool to indicate whether we have
-    # allowed_token_ids.
     # `allowed_token_ids_mask` is a 2D bool tensor of shape (max batch size,
     # vocab size).
-    no_allowed_token_ids: bool
     allowed_token_ids_mask: Optional[torch.Tensor]

vllm/v1/sample/sampler.py

@@ -237,9 +237,7 @@ def apply_allowed_token_ids(
         logits: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> torch.Tensor:
-        # One idea is implement this as a PyTorch C++ op, and we may
-        # even optimize the logit_bias layout.
-        if not sampling_metadata.no_allowed_token_ids:
+        if sampling_metadata.allowed_token_ids_mask is not None:
             logits.masked_fill_(sampling_metadata.allowed_token_ids_mask,
                                 float("-inf"))
         return logits

vllm/v1/worker/gpu_input_batch.py

@@ -192,7 +192,7 @@ def __init__(
 
         self.logit_bias: List[Optional[Dict[int,
                                             float]]] = [None] * max_num_reqs
-        self.has_allowed_token_ids: List[bool] = [False] * max_num_reqs
+        self.has_allowed_token_ids: Set[str] = set()
         self.allowed_token_ids_mask: Optional[torch.Tensor] = None
         self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
 
@@ -297,7 +297,7 @@ def add_request(
                        for tid in sampling_params.allowed_token_ids):
                 raise ValueError(
                     "allowed_token_ids contains out-of-vocab token id")
-            self.has_allowed_token_ids[req_index] = True
+            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,
@@ -357,7 +357,7 @@ def remove_request(self, req_id: str) -> Optional[int]:
             self.request_lora_mapping[req_index] = 0
 
         self.logit_bias[req_index] = None
-        self.has_allowed_token_ids[req_index] = False
+        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)
         return req_index
@@ -428,8 +428,6 @@ def condense(self, empty_req_indices: List[int]) -> None:
 
             self.logit_bias[empty_index] = self.logit_bias[last_req_index]
 
-            self.has_allowed_token_ids[
-                empty_index] = self.has_allowed_token_ids[last_req_index]
             if self.allowed_token_ids_mask_cpu_tensor is not None:
                 self.allowed_token_ids_mask_cpu_tensor[
                     empty_index] = self.allowed_token_ids_mask_cpu_tensor[
@@ -478,8 +476,8 @@ def _make_sampling_metadata(self) -> SamplingMetadata:
             prompt_token_ids = None
 
         allowed_token_ids_mask: Optional[torch.Tensor] = None
-        if not self.no_allowed_token_ids and \
-                self.allowed_token_ids_mask is not None:
+        if not self.no_allowed_token_ids:
+            assert self.allowed_token_ids_mask is not None
             copy_slice(self.allowed_token_ids_mask_cpu_tensor,
                        self.allowed_token_ids_mask, num_reqs)
             allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs]
@@ -502,7 +500,6 @@ def _make_sampling_metadata(self) -> SamplingMetadata:
             min_tokens=self.min_tokens,
             no_penalties=self.no_penalties,
             logit_bias=self.logit_bias[:num_reqs],
-            no_allowed_token_ids=self.no_allowed_token_ids,
             allowed_token_ids_mask=allowed_token_ids_mask,
         )
 
@@ -597,4 +594,4 @@ def no_prompt_logprob(self) -> bool:
 
     @property
     def no_allowed_token_ids(self) -> bool:
-        return not any(self.has_allowed_token_ids)
+        return len(self.has_allowed_token_ids) == 0