Initial commit

tests/samplers/test_sampler.py

@@ -7,7 +7,7 @@
 import torch
 from transformers import GenerationConfig, GenerationMixin
 
-from vllm.model_executor.layers.sampler import Sampler
+from vllm.model_executor.layers.sampler import ApplyToppTopkScalar, Sampler
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.model_executor.utils import set_random_seed
 from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
@@ -700,3 +700,62 @@ def test_sampling_params(sampling_params: List[SamplingParams]):
 
  assert tokens1[0] == tokens2[1]
  assert tokens1[1] == tokens2[0]
+
+
+def test_topk_topk_scalar():
+ obj1 = ApplyToppTopkScalar(2)
+ assert ApplyToppTopkScalar._padded_k == 0
+ x = torch.tensor([[9, 9, 8, 8, 8, 8, 7, 7, 7.0],
+ [10, 10, 9, 9, 9, 8, 5, 5, 5]])
+
+ retval1 = obj1(x, p=0.9, k=5)
+ ninf = -float("inf")
+ expected1 = torch.tensor([[9., 9., 8., 8., 8., 8., ninf, ninf, ninf],
+ [10., 10., 9., 9., 9., ninf, ninf, ninf, ninf]])
+ assert torch.all(retval1 == expected1).item()
+ assert ApplyToppTopkScalar._padded_k == 9
+
+ obj2 = ApplyToppTopkScalar(2)
+ assert obj2._padded_k == 9
+
+ x = torch.tensor([[2, 2, 9, 9, 2, 2, 1, 1, 1.0],
+ [10, 9, 9, 5, 9, 9, 5, 9, 10]])
+ retval2 = obj2(x, p=0.9, k=5)
+ expected2 = torch.tensor(
+ [[ninf, ninf, 9., 9., ninf, ninf, ninf, ninf, ninf],
+ [10., ninf, 9., ninf, 9., 9., ninf, 9., 10.]])
+ assert torch.all(retval2 == expected2).item()
+ assert obj2._padded_k == 9
+
+ retval3 = obj2(x, p=1.0, k=5)
+ expected3 = torch.tensor([[2., 2., 9., 9., 2., 2., ninf, ninf, ninf],
+ [10., 9., 9., ninf, 9., 9., ninf, 9., 10.]])
+
+ assert torch.all(retval3 == expected3).item()
+
+ # this should not be done in general, doing it here for testing purposes
+ ApplyToppTopkScalar._padded_k = 0
+ x = torch.tensor([[1, 1, 1, 9, 8, 1, 1, 1, 1.0],
+ [2, 1, 2, 2, 1, 1, 1, 1, 1]])
+ obj3 = ApplyToppTopkScalar(2)
+ retval4 = obj3(x, p=0.9, k=2)
+ expected4 = torch.tensor(
+ [[ninf, ninf, ninf, 9., 8., ninf, ninf, ninf, ninf],
+ [2., ninf, 2., 2., ninf, ninf, ninf, ninf, ninf]])
+ assert torch.all(retval4 == expected4).item()
+ assert obj3._padded_k == 4
+ y = torch.tensor([[8, 8, 8, 9, 8, 1, 1, 1, 1.0],
+ [2, 1, 2, 2, 1, 1, 1, 1, 1]])
+ retval5 = obj3(y, p=0.9, k=2)
+ assert obj3._padded_k == 8
+ expected5 = torch.tensor([[8., 8., 8., 9., 8., ninf, ninf, ninf, ninf],
+ [2., ninf, 2., 2., ninf, ninf, ninf, ninf,
+ ninf]])
+ assert torch.all(retval5 == expected5).item()
+ y = torch.tensor([[8, 8, 8, 9, 8, 8, 1, 1, 1.0],
+ [2, 1, 2, 2, 3, 1, 1, 1, 1]])
+ retval6 = obj3(y, p=0.9, k=2)
+ expected6 = torch.tensor([[8., 8., 8., 9., 8., 8., ninf, ninf, ninf],
+ [2., ninf, 2., 2., 3., ninf, ninf, ninf, ninf]])
+ assert torch.all(retval6 == expected6).item()
+ assert obj3._padded_k == 8

vllm/model_executor/layers/sampler.py

@@ -1,5 +1,6 @@
 """A layer that samples the next tokens from the model's outputs."""
 import itertools
+import math
 from math import inf
 from typing import Dict, List, Optional, Tuple
 
@@ -77,6 +78,13 @@ def _init_sampling_tensors(
  self._do_penalties = do_penalties
  self._do_top_p_top_k = do_top_p_top_k
  self._do_min_p = do_min_p
+ self._top_p_scalar = sampling_tensors.top_ps[0].item()
+ self._top_k_scalar = sampling_tensors.top_ks[0].item()
+ scalar_p = torch.all(sampling_tensors.top_ps == self._top_p_scalar)
+ scalar_k = torch.all(sampling_tensors.top_ks == self._top_k_scalar)
+ self._scalar_p_and_k = (scalar_p and scalar_k).item()
+ if self._scalar_p_and_k and self._do_top_p_top_k:
+ self._apply_top_k_top_p_opt = ApplyToppTopkScalar(5)
 
  def forward(
  self,
@@ -122,8 +130,13 @@ def forward(
  logits.div_(sampling_tensors.temperatures.unsqueeze(dim=1))
 
  if do_top_p_top_k:
- logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps,
- sampling_tensors.top_ks)
+ if self._scalar_p_and_k:
+ logits = self._apply_top_k_top_p_opt(logits,
+ self._top_p_scalar,
+ self._top_k_scalar)
+ else:
+ logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps,
+ sampling_tensors.top_ks)
 
  if do_min_p:
  logits = _apply_min_p(logits, sampling_tensors.min_ps)
@@ -198,6 +211,101 @@ def _get_bin_counts_and_mask(
  return bin_counts, mask
 
 
+class ApplyToppTopkScalar():
+ """
+ The original implementation of _apply_top_k_top_p is more general
+ as it uses vector topp, topk
+ However in a lot of cases, topp and topk is same for all batch elements
+ For such "scalar" topp, topk cases, we can use this class
+
+ The main optimizations in this class is:
+ Use topk instead of sort, which is much faster especially for small k.
+ However just using topk might not suffice in cases as shown below
+ Consider a tensor: 9 9 8 8 8 8 7 7 7
+ Topk, with k=5, on this yields 9 9 8 8 8
+ The value "8" is on the boundary, hence the last "8" gets snipped off
+ However the original implementation accepts all the "8"s,
+ so it should output:
+ 9 9 8 8 8 8 (6 values, even though k=5)
+ To ensure these semantics, we perform topk with _padded_k elements
+ If we find more boundary elements left over,
+ then we keep incrementing _padded_k
+ and in future calls use the expanded value of __padded_k
+
+ The increments to _padded_k should be done
+ with value > 1 to prevent excessive recompilations
+ due to dynamic shapes (the output shape of the topk)
+
+ The main logic of this is in __call__
+ This is a class instead of a function, just to keep track of
+ the monotonic non-decreasing state _padded_k
+ """
+ _padded_k = 0
+
+ def __init__(self, increment: int):
+ self._increment = increment
+
+ def __call__(self, logits: torch.Tensor, p: float, k: int):
+ if k > ApplyToppTopkScalar._padded_k:
+ ApplyToppTopkScalar._padded_k = min(k + self._increment,
+ logits.shape[1])
+
+ vals, idx = torch.topk(logits, k=ApplyToppTopkScalar._padded_k, \
+ dim=1, sorted=True)
+
+ # this "if" checks if we have bucketed so much that
+ # we have padded k upto shape of logits
+ if ApplyToppTopkScalar._padded_k != logits.shape[1]:
+ smallest_of_top_k = vals[:, k - 1]
+ num_duplicates_of_smallest_of_topk = torch.sum(
+ logits == smallest_of_top_k.unsqueeze(1), 1)
+ max_num_duplicates_of_smallest_of_topk = torch.max(
+ num_duplicates_of_smallest_of_topk).item()
+
+ # there are n repeats for a border
+ # (border meaning the smallest value of the top k).
+ # we do not know if only 1 or 2 or (n-1)
+ # of them lie outside the kth border,
+ # so we choose to conservatively increase by n-1
+ # when num_duplicates > _padded_k - k
+ if max_num_duplicates_of_smallest_of_topk - 1 > (
+ ApplyToppTopkScalar._padded_k - k):
+ incr = int(
+ math.ceil((max_num_duplicates_of_smallest_of_topk - 1) /
+ self._increment) * self._increment)
+ # this while loop should be traversed at most twice,
+ # because we dont increment by self._increment and retry
+ # instead we compute incr in one go
+ ApplyToppTopkScalar._padded_k = min(
+ ApplyToppTopkScalar._padded_k + incr, logits.shape[1])
+
+ # recompute topk with expanded padded_k
+ vals, idx = torch.topk(logits, \
+ k=ApplyToppTopkScalar._padded_k, \
+ dim=1, sorted=True)
+
+ idx = torch.fliplr(idx)
+ vals = torch.fliplr(vals)
+
+ top_k_smallest_val_idx = vals.size(1) - k
+ top_k_mask = vals[:, top_k_smallest_val_idx].unsqueeze(1)
+ top_k_mask = vals < top_k_mask
+ vals.masked_fill_(top_k_mask, -float("inf"))
+
+ probs_sort = vals.softmax(dim=-1)
+ probs_sum = probs_sort.cumsum(dim=-1)
+ top_p_mask = probs_sum <= (1 - p)
+ top_p_mask[:, -1] = False
+ vals.masked_fill_(top_p_mask, -float("inf"))
+
+ new_logits = torch.full(logits.shape,
+ -float("inf"),
+ device=logits.device)
+ new_logits.scatter_(1, idx, vals)
+
+ return new_logits
+
+
 def _apply_min_tokens_penalty(
  logits: torch.Tensor,
  sampling_metadata: SamplingMetadata,