397 lines
14 KiB
Python
397 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import random
|
|
from dataclasses import dataclass
|
|
from typing import Optional, Union
|
|
|
|
import torch
|
|
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
|
|
|
|
from vllm.engine.arg_utils import EngineArgs
|
|
from vllm.transformers_utils.tokenizer_group.base_tokenizer_group import (
|
|
BaseTokenizerGroup)
|
|
from vllm.v1.engine import EngineCoreOutput, FinishReason
|
|
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
|
|
|
|
GeneralTokenizerType = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
|
|
|
|
# Number of sample logprobs to request when testing sample logprobs
|
|
NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5
|
|
# Number of prompt logprobs to request when testing prompt logprobs
|
|
NUM_PROMPT_LOGPROBS_UNDER_TEST = 7
|
|
|
|
TOKENIZER_NAME = "meta-llama/Llama-3.2-1B"
|
|
|
|
FULL_STRINGS = [
|
|
"My name is Robert from Neural Magic and I love working on vLLM so much!",
|
|
"Red Hat is the best open source company by far across Linux, K8s, and AI.",
|
|
"Nick is the name of my brother in addition to my colleague from Red Hat.",
|
|
]
|
|
STOP_STRINGS = ["I love working on", "company by far", "brother in"]
|
|
PROMPT_LEN = 5
|
|
|
|
random.seed(42)
|
|
|
|
|
|
def _create_random_top_logprob_test_vector(
|
|
num_logprobs: int,
|
|
lower: float,
|
|
upper: float,
|
|
) -> torch.Tensor:
|
|
"""Create a random vector of top logprob float values.
|
|
|
|
Use to create fake sample logprobs for testing.
|
|
|
|
Note that a real production scenario would require
|
|
logprobs to be sorted in descending order, something
|
|
which is omitted in this function.
|
|
|
|
Args:
|
|
num_logprobs: number of top logprobs
|
|
lower: lower range of logprob float values
|
|
upper: upper range of logprob float values
|
|
|
|
Returns:
|
|
1D length-`num_logprobs` torch Tensor of float logprob values
|
|
"""
|
|
return torch.rand(num_logprobs) * (upper - lower) + lower
|
|
|
|
|
|
def _create_random_top_logprob_test_matrix(
|
|
shape: tuple,
|
|
lower: float,
|
|
upper: float,
|
|
) -> torch.Tensor:
|
|
"""Create a random matrix of top logprob float values.
|
|
|
|
Use to create fake prompt logprobs for testing.
|
|
|
|
Note that a real production scenario would require
|
|
logprobs to be sorted in descending order along rows,
|
|
something which is omitted in this function.
|
|
|
|
Args:
|
|
shape: (num_tokens,num_logprobs) tuple representing
|
|
matrix shape
|
|
lower: lower range of logprob float values
|
|
upper: upper range of logprob float values
|
|
|
|
Returns:
|
|
2D num_tokens x num_logprobs torch Tensor of float logprob values
|
|
"""
|
|
return torch.rand(*shape) * (upper - lower) + lower
|
|
|
|
|
|
def _create_random_top_token_test_vector(
|
|
num_logprobs: int,
|
|
lower: int,
|
|
upper: int,
|
|
sampled_token_id: int,
|
|
adjust_num_logprobs: bool = True) -> tuple[torch.Tensor, int]:
|
|
"""Create a random vector of top logprob token indices
|
|
|
|
Use to create fake sample logprobs for testing. The sampled token
|
|
ID must always be one of the top logprobs, which this dummy test
|
|
vector generator enforces. OpenAI API
|
|
compatible engines must be able to return an additional sample
|
|
logprob for the sampled token if the sampled token was not
|
|
among the top sample logprobs; `adjust_num_logprobs` emulates
|
|
this behavior by increasing the vector length by 1 if
|
|
`adjust_num_logprobs` is set.
|
|
|
|
Args:
|
|
num_logprobs: number of top logprobs
|
|
lower: lower range of token ids
|
|
upper: upper range of token ids
|
|
sampled_token_id: the token actually sampled
|
|
adjust_num_logprobs: if True, emulate situation where sampled
|
|
token logprob must be injected into top
|
|
logprobs
|
|
|
|
Returns:
|
|
1D length-x torch Tensor of token ids where x is
|
|
`num_logprobs+1` if `adjust_num_logprobs` and
|
|
`num_logprobs` otherwise
|
|
sampled_token_rank: the rank of sampled_token_id in the vocab
|
|
vector when sorted in descending order by
|
|
logprob
|
|
"""
|
|
|
|
# Calculate the final number of logprobs required
|
|
total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs
|
|
|
|
# Generate random indices using torch
|
|
choice_tensor = torch.randperm(upper - lower)[:total_logprobs] + lower
|
|
|
|
# Ensure the sampled token ID is included in the tensor
|
|
choice_tensor[0] = sampled_token_id
|
|
|
|
# Check if the sampled_token_id occurs in choice_tensor[1:]
|
|
if sampled_token_id in choice_tensor[1:]:
|
|
sampled_token_rank = (choice_tensor[1:] == sampled_token_id).nonzero(
|
|
as_tuple=True)[0].item()
|
|
else:
|
|
# If not found, assign a random int between num_logprobs and 50700
|
|
sampled_token_rank = random.randint(num_logprobs, 50700)
|
|
|
|
return choice_tensor, sampled_token_rank
|
|
|
|
|
|
def _create_random_top_token_test_matrix(
|
|
shape: tuple[int, int],
|
|
lower: int,
|
|
upper: int,
|
|
tokens_list: list[int],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Create a random matrix of top logprob token indices
|
|
|
|
Use to create fake prompt logprobs for testing.
|
|
|
|
Token ids are generated randomly and sampled without
|
|
replacement.
|
|
|
|
Args:
|
|
shape: (num_tokens, num_logprobs) tuple representing
|
|
matrix shape
|
|
lower: lower range of token ids
|
|
upper: upper range of token ids
|
|
|
|
Returns:
|
|
tuple containing:
|
|
- 2D num_tokens x num_logprobs+1 torch Tensor of token ids
|
|
- 1D tensor of ranks of prompt tokens in their respective
|
|
rows, or random values
|
|
"""
|
|
num_elements = shape[0] * shape[1]
|
|
choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower
|
|
matrix = torch.cat(
|
|
(torch.tensor(tokens_list, dtype=torch.int).unsqueeze(-1),
|
|
choice_tensor.view(shape)),
|
|
dim=1)
|
|
|
|
# Initialize the tensor for storing the ranks
|
|
prompt_token_ranks = torch.empty(shape[0], dtype=torch.int)
|
|
|
|
# Iterate over each row to check presence of
|
|
# tokens_list[rdx] and determine its index
|
|
for rdx in range(shape[0]):
|
|
row = matrix[rdx,
|
|
1:] # Skip the first column as it contains the token list
|
|
token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0]
|
|
if token_index.numel() > 0:
|
|
prompt_token_ranks[rdx] = token_index.item()
|
|
else:
|
|
prompt_token_ranks[rdx] = random.randint(shape[1], 50700)
|
|
|
|
return matrix, prompt_token_ranks
|
|
|
|
|
|
def decode_token(
|
|
tok_id: int,
|
|
tokenizer: PreTrainedTokenizer,
|
|
) -> str:
|
|
"""Reproduce the process of detokenizing a token for testing purposes.
|
|
|
|
Args:
|
|
tok_id: token id to detokenize
|
|
tokenizer: tokenizer to use for detokenization
|
|
|
|
Returns:
|
|
string representation of token
|
|
"""
|
|
return tokenizer.convert_ids_to_tokens(tok_id)
|
|
|
|
|
|
def generate_dummy_sample_logprobs(
|
|
sampled_tokens_list: list,
|
|
num_logprobs: int,
|
|
tokenizer: PreTrainedTokenizer,
|
|
) -> list[tuple[list[int], list[float], int]]:
|
|
"""Generate dummy sample logprobs
|
|
|
|
Generate a test data structure which imitates the list of sample logprobs
|
|
which would be assembled in the engine core during decode phase.
|
|
|
|
Args:
|
|
sampled_tokens_list: list of sampled tokens
|
|
num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token
|
|
tokenizer: model tokenizer to use for detokenization
|
|
|
|
Returns
|
|
list of (top token ids vector, logprobs vector, sampled token rank)
|
|
Python lists tuples; in each tuple the logprobs and top token ids
|
|
vectors have the same length which is either `num_logprobs` or
|
|
`num_logprobs+1`. Sampled token rank is the rank (index+1) of the
|
|
sampled token within the vocab vector when sorted by logprob in
|
|
descending order.
|
|
"""
|
|
res = []
|
|
for sampled_token_id in sampled_tokens_list:
|
|
(
|
|
token_vector,
|
|
sampled_token_rank,
|
|
) = _create_random_top_token_test_vector(num_logprobs, 0,
|
|
len(tokenizer.vocab) - 1,
|
|
sampled_token_id)
|
|
|
|
res.append(
|
|
(token_vector,
|
|
_create_random_top_logprob_test_vector(num_logprobs + 1, -100,
|
|
0), sampled_token_rank))
|
|
|
|
# Convert tensors in the list tuples to Python lists
|
|
res_list_format = [
|
|
(log_probs_tensor.tolist(), token_ids_tensor.tolist(),
|
|
sampled_token_rank)
|
|
for log_probs_tensor, token_ids_tensor, sampled_token_rank in res
|
|
]
|
|
|
|
return res_list_format
|
|
|
|
|
|
def generate_dummy_prompt_logprobs_tensors(
|
|
prompt_tokens_list: list,
|
|
num_logprobs: int,
|
|
tokenizer: PreTrainedTokenizer,
|
|
) -> LogprobsTensors:
|
|
"""Generate dummy prompt logprobs tensors
|
|
|
|
Generate a test data structure which imitates the torch Tensors of prompt
|
|
logprobs which would be assembled in the engine core during chunked
|
|
prefill.
|
|
|
|
Args:
|
|
prompt_tokens_list: list of prompt tokens
|
|
num_logprobs: return `num_logprobs` logprobs per token
|
|
tokenizer: model tokenizer to use for detokenization
|
|
|
|
Returns
|
|
Single tuple of (logprobs matrix, top token ids matrix) torch Tensor,
|
|
where both matrices have dimensions
|
|
num_prompt_tokens x num_logprobs
|
|
"""
|
|
# For now, assume the whole prompt is processed in one chunk; thus,
|
|
# the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`.
|
|
# Prior to injecting `None` at the beginning of prompt logprobs (which
|
|
# happens later in the detokenizer, not here), the prompt logprobs in
|
|
# the ith position are predicting the probability distribution of the
|
|
# prompt token in (i+1)st position. Thus, we concat
|
|
# `prompt_tokens_list[1:]` to the dummy token ids, just as the engine
|
|
# would.
|
|
num_prompt_logprobs = len(prompt_tokens_list) - 1
|
|
(
|
|
token_vector,
|
|
prompt_token_ranks,
|
|
) = _create_random_top_token_test_matrix(
|
|
(num_prompt_logprobs, num_logprobs), 0,
|
|
len(tokenizer.vocab) - 1, prompt_tokens_list[1:])
|
|
return LogprobsTensors(
|
|
token_vector,
|
|
_create_random_top_logprob_test_matrix(
|
|
(num_prompt_logprobs, num_logprobs + 1), -100, 0),
|
|
prompt_token_ranks)
|
|
|
|
|
|
@dataclass
|
|
class DummyOutputProcessorTestVectors:
|
|
"""Dummy test vectors for output processor tests"""
|
|
tokenizer: GeneralTokenizerType
|
|
tokenizer_group: BaseTokenizerGroup
|
|
vllm_config: EngineArgs
|
|
full_tokens: list[list[int]] # Prompt + generated tokens
|
|
prompt_tokens: list[list[int]]
|
|
generation_tokens: list[list[int]]
|
|
# Each request is associated with a tuple of
|
|
# (top tokens, top logprobs, ranks) prompt logprobs tensors
|
|
prompt_logprobs: list[LogprobsTensors]
|
|
# Each request is associated with a sample logprobs; a request's
|
|
# sample logprobs are a list of (top tokens, top logprobs, ranks)
|
|
# sample logprobs tensors at each sequence position
|
|
generation_logprobs: list[list[tuple[list[int], list[float], int]]]
|
|
prompt_strings: list[str]
|
|
prompt_strings_len: list[int]
|
|
generation_strings: list[str]
|
|
|
|
|
|
class MockEngineCore:
|
|
"""Mock engine core outputs form premade tokens lists."""
|
|
|
|
def __init__(
|
|
self,
|
|
tokens_list: list[list[int]],
|
|
# For each request, for each sampled token offset,
|
|
# a tuple of
|
|
# (list of topk token ids, list of sample logprob vals, rank)
|
|
generated_logprobs_raw: Optional[list[list[tuple[list[int],
|
|
list[float],
|
|
int]]]] = None,
|
|
# For each request, a tuple of
|
|
# (prompt logprob val matrix, prompt logprob tok id matrix);
|
|
# each matrix has dimensions
|
|
# (num prompt toks) x (num prompt logprobs+1)
|
|
prompt_logprobs_raw: Optional[list[LogprobsTensors]] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
stop_token_ids: Optional[list[int]] = None,
|
|
ignore_eos: bool = False,
|
|
) -> None:
|
|
self.num_requests = len(tokens_list)
|
|
self.tokens_list = tokens_list
|
|
self.current_idx = 0
|
|
self.generated_logprobs_raw = generated_logprobs_raw
|
|
self.do_logprobs = generated_logprobs_raw is not None
|
|
self.prompt_logprobs_raw = prompt_logprobs_raw
|
|
self.do_prompt_logprobs = prompt_logprobs_raw is not None
|
|
self.request_finished = [False for _ in range(self.num_requests)]
|
|
self.eos_token_id = eos_token_id
|
|
self.stop_token_ids = stop_token_ids
|
|
self.ignore_eos = ignore_eos
|
|
|
|
def get_outputs(self) -> list[EngineCoreOutput]:
|
|
do_logprobs = self.do_logprobs
|
|
do_prompt_logprobs = self.do_prompt_logprobs
|
|
token_idx = self.current_idx
|
|
|
|
outputs = []
|
|
for req_idx, token_ids in enumerate(self.tokens_list):
|
|
if not self.request_finished[req_idx]:
|
|
if do_logprobs:
|
|
assert self.generated_logprobs_raw is not None
|
|
(logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (
|
|
self.generated_logprobs_raw[req_idx][token_idx])
|
|
logprobs = LogprobsLists(
|
|
[logprobs_token_ids_],
|
|
[logprobs_],
|
|
[sampled_token_ranks_],
|
|
)
|
|
else:
|
|
logprobs = None
|
|
if do_prompt_logprobs:
|
|
if self.current_idx == 0:
|
|
assert self.prompt_logprobs_raw is not None
|
|
prompt_logprobs = self.prompt_logprobs_raw[req_idx]
|
|
else:
|
|
prompt_logprobs = None
|
|
else:
|
|
prompt_logprobs = None
|
|
new_token_id = token_ids[token_idx]
|
|
output = EngineCoreOutput(
|
|
request_id=f"request-{req_idx}",
|
|
new_token_ids=[new_token_id],
|
|
new_logprobs=logprobs,
|
|
new_prompt_logprobs_tensors=prompt_logprobs,
|
|
)
|
|
if token_idx == len(token_ids) - 1:
|
|
output.finish_reason = FinishReason.LENGTH
|
|
self.request_finished[req_idx] = True
|
|
if not self.ignore_eos and new_token_id == self.eos_token_id:
|
|
output.finish_reason = FinishReason.STOP
|
|
self.request_finished[req_idx] = True
|
|
if new_token_id in (self.stop_token_ids or ()):
|
|
output.finish_reason = FinishReason.STOP
|
|
output.stop_reason = new_token_id
|
|
self.request_finished[req_idx] = True
|
|
outputs.append(output)
|
|
|
|
self.current_idx += 1
|
|
return outputs
|