[RFC][V1] LogitsProcessor interface

Signed-off-by: Nick Hill <[email protected]>

Author: njhill

vllm/v1/sample/logits_processor.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+import dataclasses
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Sequence, Set, Tuple
+
+import torch
+
+from vllm import SamplingParams
+
+
+@dataclasses.dataclass
+class BatchUpdate:
+    # Batch indices of any removed requests.
+    removed: List[int]
+    # (from, to) batch indices of any requests
+    # moved within the batch.
+    moved: List[Tuple[int, int]]
+    # (index, params, output_tok_ids) for new
+    # requests added to the batch.
+    added: List[Tuple[int, SamplingParams, List[int]]]
+    # The current number of requests in the batch.
+    batch_size: int
+
+
+class LogitsProcessor(ABC):
+
+    @abstractmethod
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_states(
+        self,
+        batch_update: Optional[BatchUpdate] = None,
+    ) -> None:
+        """Called when there are new output tokens, prior
+        to each forward pass.
+
+        Args:
+            batch_update is non-None iff there have been
+            changes to the batch makeup.
+        """
+        raise NotImplementedError
+
+
+###### ----- LogitsProcessor impls below here
+
+
+class MinPLogitsProcessor(LogitsProcessor):
+
+    def __init__(self, max_num_reqs: int, pin_memory: bool,
+                 device: torch.device):
+        self.min_p_count: int = 0
+
+        self.min_p_cpu_tensor = torch.zeros((max_num_reqs, ),
+                                            dtype=torch.float32,
+                                            device="cpu",
+                                            pin_memory=pin_memory)
+        self.min_p_cpu = self.min_p_cpu_tensor.numpy()
+        # Pre-allocated device tensor
+        self.min_p_gpu: torch.Tensor = torch.empty((max_num_reqs, ),
+                                                   dtype=torch.float32,
+                                                   device=device)
+        # Current slice of the device tensor
+        self.min_p: torch.Tensor = self.min_p_gpu[:0]
+
+    def update_states(self, batch_update: Optional[BatchUpdate] = None):
+        if not batch_update:
+            return
+
+        needs_update = False
+        if self.min_p_count:
+            # Process removed and moved requests.
+            for index in batch_update.removed:
+                if self.min_p_cpu[index]:
+                    self.min_p_count -= 1
+                    needs_update = True
+
+            for from_index, to_index in batch_update.moved:
+                min_p = self.min_p_cpu[from_index]
+                self.min_p_cpu[to_index] = min_p
+                if min_p:
+                    needs_update = True
+
+        # Process added requests.
+        for index, sampling_params, _ in batch_update.added:
+            min_p = sampling_params.min_p
+            self.min_p_cpu[index] = min_p
+            if min_p:
+                self.min_p_count += 1
+                needs_update = True
+
+        # Update tensors if needed.
+        size = batch_update.batch_size
+        if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
+
+            self.min_p = self.min_p_gpu[:size]
+            self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
+            self.min_p.unsqueeze_(1)
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if not self.min_p_count:
+            return logits
+
+        # Convert logits to probability distribution
+        probability_values = torch.nn.functional.softmax(logits, dim=-1)
+        # Calculate maximum probabilities per sequence
+        max_probabilities = torch.amax(probability_values,
+                                       dim=-1,
+                                       keepdim=True)
+        # Adjust min_p
+        adjusted_min_p = max_probabilities.mul_(self.min_p)
+        # Identify valid tokens using threshold comparison
+        invalid_token_mask = probability_values < adjusted_min_p
+        # Apply mask using boolean indexing
+        logits[invalid_token_mask] = -float('inf')
+        return logits
+
+
+class LogitBiasLogitsProcessor(LogitsProcessor):
+
+    def __init__(self, pin_memory: bool, device: torch.device):
+        self.biases: Dict[int, Dict[int, float]] = {}
+        self.device = device
+        self.pin_memory = pin_memory
+
+        self.bias_tensor: torch.Tensor = torch.tensor(())
+        self.logits_slice: Tuple[torch.Tensor, torch.Tensor] = (torch.tensor(
+            ()), torch.tensor(()))
+
+    def update_states(self, batch_update: Optional[BatchUpdate] = None):
+        if not batch_update:
+            return
+
+        needs_update = False
+        if self.biases:
+            # Process removed and moved requests.
+            for index in batch_update.removed:
+                if self.biases.pop(index, None):
+                    needs_update = True
+
+            for from_index, to_index in batch_update.moved:
+                if entry := self.biases.pop(from_index, None):
+                    self.biases[to_index] = entry
+                    needs_update = True
+
+        # Process added requests.
+        for index, sampling_params, _ in batch_update.added:
+            if lb := sampling_params.logit_bias:
+                self.biases[index] = lb
+                needs_update = True
+
+        # Update tensors if needed.
+        if self.biases and needs_update:
+            reqs, tok_ids, biases = [], [], []
+            for req, lb in self.biases.items():
+                reqs.extend([req] * len(lb))
+                tok_ids.extend(lb.keys())
+                biases.extend(lb.values())
+
+            self.bias_tensor = self._tensor(biases, torch.float32)
+            self.logits_slice = (self._tensor(reqs, torch.int32),
+                                 self._tensor(tok_ids, torch.int32))
+
+    def _tensor(self, data: List, dtype: torch.dtype) -> torch.Tensor:
+        return (torch.tensor(data,
+                             device="cpu",
+                             dtype=dtype,
+                             pin_memory=self.pin_memory).to(device=self.device,
+                                                            non_blocking=True))
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.biases:
+            logits[self.logits_slice] += self.bias_tensor
+        return logits
+
+
+class MinTokensLogitsProcessor(LogitsProcessor):
+
+    def __init__(self, pin_memory: bool, device: torch.device):
+        # index -> (min_toks, output_token_ids, stop_token_ids)
+        self.min_toks: Dict[int, Tuple[int, Sequence[int], Set[int]]] = {}
+        self.device = device
+        self.pin_memory = pin_memory
+
+        self.logits_slice: Tuple[torch.Tensor, torch.Tensor] = (torch.tensor(
+            ()), torch.tensor(()))
+
+    def update_states(self, batch_update: Optional[BatchUpdate] = None):
+        needs_update = False
+        if batch_update:
+            if self.min_toks:
+                # Process removed and moved requests.
+                for index in batch_update.removed:
+                    if self.min_toks.pop(index, None):
+                        needs_update = True
+
+                for from_index, to_index in batch_update.moved:
+                    if entry := self.min_toks.pop(from_index, None):
+                        self.min_toks[to_index] = entry
+                        needs_update = True
+
+            # Process added requests.
+            for index, sampling_params, output_tok_ids in batch_update.added:
+                if ((min_tokens := sampling_params.min_tokens)
+                        and len(output_tok_ids) < min_tokens):
+                    self.min_toks[index] = (min_tokens, output_tok_ids,
+                                            sampling_params.all_stop_token_ids)
+                    needs_update = True
+
+        if self.min_toks:
+            # Check for any requests that have attained their min tokens.
+            to_remove = tuple(index for index, (min_toks, out_tok_ids,
+                                                _) in self.min_toks.items()
+                              if len(out_tok_ids) >= min_toks)
+            if to_remove:
+                needs_update = True
+                for index in to_remove:
+                    del self.min_toks[index]
+
+            # Update tensors if needed.
+            if needs_update and self.min_toks:
+                reqs: List[int] = []
+                tok_ids: List[int] = []
+                for req, (_, _, stop_tok_ids) in self.min_toks.items():
+                    reqs.extend([req] * len(stop_tok_ids))
+                    tok_ids.extend(stop_tok_ids)
+
+                self.logits_slice = (self._tensor(reqs, torch.int32),
+                                     self._tensor(tok_ids, torch.int32))
+
+    def _tensor(self, data: List, dtype: torch.dtype) -> torch.Tensor:
+        return (torch.tensor(data,
+                             device="cpu",
+                             dtype=dtype,
+                             pin_memory=self.pin_memory).to(device=self.device,
+                                                            non_blocking=True))
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.min_toks:
+            logits[self.logits_slice] = -float("inf")
+        return logits

vllm/v1/sample/metadata.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional
 
 import torch
 
+from vllm.v1.sample.logits_processor import LogitsProcessor
+
 
 @dataclass
 class SamplingMetadata:
@@ -18,7 +20,6 @@ class SamplingMetadata:
 
     top_p: Optional[torch.Tensor]
     top_k: Optional[torch.Tensor]
-    min_p: Optional[torch.Tensor]
 
     generators: Dict[int, torch.Generator]
 
@@ -33,7 +34,5 @@ class SamplingMetadata:
 
     output_token_ids: List[List[int]]
 
-    # req_index -> (min_tokens, stop_token_ids)
-    min_tokens: Dict[int, Tuple[int, Set[int]]]
-
-    logit_bias: List[Optional[Dict[int, float]]]
+    logits_procs: List[LogitsProcessor]
+    nongreedy_logits_procs: List[LogitsProcessor]

vllm/v1/sample/ops/penalties.py

@@ -1,29 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Dict, List, Set, Tuple
+from typing import List
 
 import torch
 
 from vllm.model_executor.layers.utils import apply_penalties
 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
 
 
-def apply_min_token_penalties(
-        logits: torch.Tensor, output_token_ids: List[List[int]],
-        min_tokens: Dict[int, Tuple[int, Set[int]]]) -> None:
-    """
-    Applies minimum token penalty by setting the logits of the stop tokens
-    to -inf.
-    """
-    min_tokens_logits_to_penalize: List[Tuple[int, int]] = []
-    for index, (min_token, stop_token_ids) in min_tokens.items():
-        if len(output_token_ids[index]) < min_token:
-            for stop_token_id in stop_token_ids:
-                min_tokens_logits_to_penalize.append((index, stop_token_id))
-    if min_tokens_logits_to_penalize:
-        logits[tuple(zip(*min_tokens_logits_to_penalize))] = -float("inf")
-
-
 def apply_all_penalties(
     logits: torch.Tensor,
     prompt_token_ids: torch.Tensor,

vllm/v1/sample/sampler.py

@@ -6,8 +6,7 @@
 
 from vllm.v1.outputs import LogprobsTensors, SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
-from vllm.v1.sample.ops.penalties import (apply_all_penalties,
-                                          apply_min_token_penalties)
+from vllm.v1.sample.ops.penalties import apply_all_penalties
 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
 from vllm.v1.sample.rejection_sampler import RejectionSampler
 
@@ -47,8 +46,11 @@ def forward(
 
         # Use float32 for the logits.
         logits = logits.to(torch.float32)
-        # Apply logits bias.
-        logits = self.apply_logits_bias(logits, sampling_metadata)
+
+        # Apply logits processors.
+        for processor in sampling_metadata.logits_procs:
+            logits = processor.apply(logits)
+
         # Apply penalties (e.g., min_tokens, freq_penalties).
         logits = self.apply_penalties(logits, sampling_metadata)
         # Sample the next token.
@@ -103,9 +105,9 @@ def sample(
         # Apply temperature.
         logits = self.apply_temperature(logits, sampling_metadata.temperature)
 
-        # Apply min_p.
-        if sampling_metadata.min_p is not None:
-            logits = self.apply_min_p(logits, sampling_metadata.min_p)
+        # Apply logits processors.
+        for processor in sampling_metadata.nongreedy_logits_procs:
+            logits = processor.apply(logits)
 
         # Apply top_k and/or top_p.
         random_sampled = self.topk_topp_sampler(
@@ -177,10 +179,6 @@ def apply_penalties(
         logits: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> torch.Tensor:
-        if sampling_metadata.min_tokens:
-            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(
@@ -190,39 +188,3 @@ def apply_penalties(
                 sampling_metadata.repetition_penalties,
                 sampling_metadata.output_token_ids)
         return logits
-
-    def apply_min_p(
-        self,
-        logits: torch.Tensor,
-        min_p: torch.Tensor,
-    ) -> torch.Tensor:
-        """
-        Filters logits using adaptive probability thresholding.
-        """
-        # Convert logits to probability distribution
-        probability_values = torch.nn.functional.softmax(logits, dim=-1)
-        # Calculate maximum probabilities per sequence
-        max_probabilities = torch.amax(probability_values,
-                                       dim=-1,
-                                       keepdim=True)
-        # Reshape min_p for broadcasting
-        adjusted_min_p = min_p.unsqueeze(1) * max_probabilities
-        # Identify valid tokens using threshold comparison
-        valid_token_mask = probability_values >= adjusted_min_p
-        # Apply mask using boolean indexing
-        logits[~valid_token_mask] = -float('inf')
-        return logits
-
-    def apply_logits_bias(
-        self,
-        logits: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> torch.Tensor:
-        # TODO(houseroad): this implementation is extremely inefficient.
-        # One idea is implement this as a PyTorch C++ op, and we may
-        # even optimize the logit_bias layout.
-        for i, logit_bias in enumerate(sampling_metadata.logit_bias):
-            if logit_bias:
-                for token_id, bias in logit_bias.items():
-                    logits[i, token_id] += bias
-        return logits