chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test that batched_count_greater_than does not trigger 0/1 specialization
|
||||
recompiles when batch_size varies."""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.sample.ops.logprobs import batched_count_greater_than
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
|
||||
def test_batched_count_greater_than_correctness():
|
||||
"""Basic correctness: counts elements >= the corresponding value."""
|
||||
x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=DEVICE)
|
||||
values = torch.tensor([[2.0], [5.0]], device=DEVICE)
|
||||
result = batched_count_greater_than(x, values)
|
||||
expected = torch.tensor([2, 2], device=DEVICE)
|
||||
torch.testing.assert_close(result, expected)
|
||||
|
||||
|
||||
def test_gather_logprobs_no_recompile():
|
||||
"""Sampler.gather_logprobs with batch_size=1 then 2 must not recompile.
|
||||
|
||||
This guards against 0/1 specialization: dynamo normally specializes on
|
||||
tensor sizes 0 and 1, causing a recompile when the size first exceeds 1.
|
||||
The mark_unbacked calls in gather_logprobs prevent this.
|
||||
"""
|
||||
torch._dynamo.reset()
|
||||
|
||||
compile_count = 0
|
||||
orig_backend = current_platform.simple_compile_backend
|
||||
|
||||
def counting_backend(gm, example_inputs):
|
||||
nonlocal compile_count
|
||||
compile_count += 1
|
||||
if orig_backend == "inductor":
|
||||
return torch._inductor.compile(gm, example_inputs)
|
||||
return gm
|
||||
|
||||
# Monkey-patch batched_count_greater_than with our counting backend
|
||||
# so we can detect recompiles through the production code path.
|
||||
import vllm.v1.sample.ops.logprobs as logprobs_module
|
||||
import vllm.v1.sample.sampler as sampler_module
|
||||
|
||||
unwrapped = batched_count_greater_than._torchdynamo_orig_callable
|
||||
patched = torch.compile(unwrapped, backend=counting_backend)
|
||||
orig_fn = logprobs_module.batched_count_greater_than
|
||||
|
||||
logprobs_module.batched_count_greater_than = patched
|
||||
sampler_module.batched_count_greater_than = patched
|
||||
|
||||
try:
|
||||
vocab_size = 32
|
||||
num_logprobs = 3
|
||||
|
||||
# Call 1: batch_size=1
|
||||
logprobs1 = torch.randn(1, vocab_size, device=DEVICE)
|
||||
token_ids1 = torch.randint(
|
||||
0, vocab_size, (1,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
Sampler.gather_logprobs(logprobs1, num_logprobs, token_ids1)
|
||||
assert compile_count == 1, f"Expected 1 compile, got {compile_count}"
|
||||
|
||||
# Call 2: batch_size=2 — should NOT recompile
|
||||
logprobs2 = torch.randn(2, vocab_size, device=DEVICE)
|
||||
token_ids2 = torch.randint(
|
||||
0, vocab_size, (2,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
Sampler.gather_logprobs(logprobs2, num_logprobs, token_ids2)
|
||||
assert compile_count == 1, (
|
||||
f"Recompiled on batch_size 1->2 (0/1 specialization). "
|
||||
f"Expected 1 compile, got {compile_count}"
|
||||
)
|
||||
|
||||
# Call 3: batch_size=8 — should NOT recompile
|
||||
logprobs3 = torch.randn(8, vocab_size, device=DEVICE)
|
||||
token_ids3 = torch.randint(
|
||||
0, vocab_size, (8,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
Sampler.gather_logprobs(logprobs3, num_logprobs, token_ids3)
|
||||
assert compile_count == 1, (
|
||||
f"Recompiled on batch_size change. Expected 1 compile, got {compile_count}"
|
||||
)
|
||||
finally:
|
||||
# Restore original function
|
||||
logprobs_module.batched_count_greater_than = orig_fn
|
||||
sampler_module.batched_count_greater_than = orig_fn
|
||||
torch._dynamo.reset()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,57 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import lm_eval
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
# arc-easy uses prompt_logprobs=1, logprobs=1
|
||||
TASK = "arc_easy"
|
||||
FILTER = "acc_norm,none"
|
||||
RTOL = 0.03
|
||||
EXPECTED_VALUE = 0.62
|
||||
|
||||
# FIXME(rob): enable prefix caching once supported.
|
||||
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
MODEL_ARGS = f"pretrained={MODEL},enforce_eager=True,enable_prefix_caching=False,gpu_memory_utilization=0.8" # noqa: E501
|
||||
SERVER_ARGS = [
|
||||
"--enforce_eager",
|
||||
"--no_enable_prefix_caching",
|
||||
"--gpu-memory-utilization=0.8",
|
||||
]
|
||||
NUM_CONCURRENT = 100
|
||||
|
||||
|
||||
def test_prompt_logprobs_e2e():
|
||||
results = lm_eval.simple_evaluate(
|
||||
model="vllm", model_args=MODEL_ARGS, tasks=TASK, batch_size="auto"
|
||||
)
|
||||
|
||||
measured_value = results["results"][TASK][FILTER]
|
||||
assert (
|
||||
measured_value - RTOL < EXPECTED_VALUE
|
||||
and measured_value + RTOL > EXPECTED_VALUE
|
||||
), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
|
||||
|
||||
|
||||
def test_prompt_logprobs_e2e_server():
|
||||
with RemoteOpenAIServer(MODEL, SERVER_ARGS) as remote_server:
|
||||
url = f"{remote_server.url_for('v1')}/completions"
|
||||
|
||||
model_args = (
|
||||
f"model={MODEL},"
|
||||
f"base_url={url},"
|
||||
f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
|
||||
)
|
||||
|
||||
results = lm_eval.simple_evaluate(
|
||||
model="local-completions",
|
||||
model_args=model_args,
|
||||
tasks=TASK,
|
||||
)
|
||||
|
||||
measured_value = results["results"][TASK][FILTER]
|
||||
assert (
|
||||
measured_value - RTOL < EXPECTED_VALUE
|
||||
and measured_value + RTOL > EXPECTED_VALUE
|
||||
), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,450 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.sample.utils import create_allowed_token_ids
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.platform_utils import is_pin_memory_available
|
||||
from vllm.utils.torch_utils import make_tensor_with_pad
|
||||
from vllm.v1.sample.logits_processor import LogitsProcessors
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
PIN_MEMORY_AVAILABLE = is_pin_memory_available()
|
||||
MAX_NUM_REQS = 256
|
||||
VOCAB_SIZE = 1024
|
||||
NUM_OUTPUT_TOKENS = 20
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
DEVICES = [
|
||||
f"{DEVICE_TYPE}:{i}"
|
||||
for i in range(1 if current_platform.device_count() == 1 else 2)
|
||||
]
|
||||
MAX_NUM_PROMPT_TOKENS = 64
|
||||
|
||||
|
||||
def _create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor:
|
||||
fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float)
|
||||
return fake_logits
|
||||
|
||||
|
||||
def _create_penalty_tensor(
|
||||
batch_size: int, penalty_value: float, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
return torch.full(
|
||||
(batch_size,), fill_value=penalty_value, dtype=torch.float, device=device
|
||||
)
|
||||
|
||||
|
||||
def _create_prompt_tokens_tensor(
|
||||
prompt_token_ids: list[list[int]],
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
return make_tensor_with_pad(
|
||||
prompt_token_ids,
|
||||
pad=vocab_size,
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
pin_memory=False,
|
||||
)
|
||||
|
||||
|
||||
def _create_bad_words_token_ids(
|
||||
batch_size: int,
|
||||
vocab_size: int,
|
||||
bad_words_lengths: tuple[int, ...],
|
||||
) -> dict[int, list[list[int]]]:
|
||||
bad_words_token_ids = {}
|
||||
for batch_idx in range(batch_size):
|
||||
token_ids_single_batch = []
|
||||
for bad_words_length in bad_words_lengths:
|
||||
token_ids = np.random.choice(
|
||||
vocab_size, size=bad_words_length, replace=True
|
||||
).tolist()
|
||||
token_ids_single_batch.append(token_ids)
|
||||
bad_words_token_ids[batch_idx] = token_ids_single_batch
|
||||
if batch_size >= 2:
|
||||
# Test no bad_words for some batch
|
||||
no_bad_words_batch_idx = np.random.choice(batch_size)
|
||||
bad_words_token_ids.pop(no_bad_words_batch_idx, None)
|
||||
return bad_words_token_ids
|
||||
|
||||
|
||||
# Returns all last tokens of bad word sequences that share the same prefix
|
||||
# as `given_prefix` (excluding the last token).
|
||||
def _collect_suffixes_with_same_prefix(
|
||||
given_prefix: list[int], bad_words_token_ids: list[list[int]]
|
||||
) -> list[int]:
|
||||
return [bwt[-1] for bwt in bad_words_token_ids if bwt[:-1] == given_prefix]
|
||||
|
||||
|
||||
# generate a valid token id that is not in bad_words_token_ids
|
||||
def _generate_valid_token_id(
|
||||
bad_words_token_ids: list[list[int]], vocab_size: int
|
||||
) -> int:
|
||||
forbidden_start_tokens = set()
|
||||
for bad_word in bad_words_token_ids:
|
||||
forbidden_start_tokens.add(bad_word[0])
|
||||
# Get a safe token that's not in forbidden starts
|
||||
safe_token_candidates = list(set(range(vocab_size)) - forbidden_start_tokens)
|
||||
# Pick a random safe token
|
||||
return np.random.choice(safe_token_candidates)
|
||||
|
||||
|
||||
def _update_output_token_ids_for_bad_words(
|
||||
metadata: SamplingMetadata, vocab_size: int
|
||||
) -> dict[int, list[int]]:
|
||||
bad_words_last_tokens = {}
|
||||
for batch_idx, bad_words_token_ids in metadata.bad_words_token_ids.items():
|
||||
output_token_ids = metadata.output_token_ids[batch_idx]
|
||||
bad_words_last_token: list[int] = []
|
||||
for i, bad_word_token_ids in enumerate(bad_words_token_ids):
|
||||
if len(bad_word_token_ids) == 1:
|
||||
# Single token id always affects logits
|
||||
bad_words_last_token.append(bad_word_token_ids[0])
|
||||
else:
|
||||
prefix_length = len(bad_word_token_ids) - 1
|
||||
has_bad_words = np.random.choice([True, False])
|
||||
if has_bad_words:
|
||||
prefix = bad_word_token_ids[:-1]
|
||||
output_token_ids[-prefix_length:] = prefix
|
||||
# Collect all last tokens from other bad words
|
||||
# that share this prefix
|
||||
bad_words_last_token.extend(
|
||||
_collect_suffixes_with_same_prefix(prefix, bad_words_token_ids)
|
||||
)
|
||||
break # Maximum one update to output_token_ids
|
||||
else: # Make sure no accidental match to bad words
|
||||
output_token_ids[-1] = _generate_valid_token_id(
|
||||
bad_words_token_ids, vocab_size
|
||||
)
|
||||
bad_words_last_tokens[batch_idx] = bad_words_last_token
|
||||
return bad_words_last_tokens
|
||||
|
||||
|
||||
def _create_default_sampling_metadata(
|
||||
num_output_tokens: int,
|
||||
batch_size: int,
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
) -> SamplingMetadata:
|
||||
output_token_ids: list[list[int]] = []
|
||||
prompt_token_ids: list[list[int]] = []
|
||||
for _ in range(batch_size):
|
||||
output_token_ids.append(
|
||||
np.random.randint(0, vocab_size, size=num_output_tokens).tolist()
|
||||
)
|
||||
prompt_token_ids.append(
|
||||
np.random.randint(
|
||||
0, vocab_size, size=np.random.randint(1, MAX_NUM_PROMPT_TOKENS)
|
||||
).tolist()
|
||||
)
|
||||
fake_sampling_metadata = SamplingMetadata(
|
||||
temperature=torch.full((batch_size,), 0.0),
|
||||
all_greedy=True,
|
||||
all_random=False,
|
||||
top_p=None,
|
||||
top_k=None,
|
||||
generators={},
|
||||
max_num_logprobs=0,
|
||||
prompt_token_ids=_create_prompt_tokens_tensor(
|
||||
prompt_token_ids, vocab_size, device
|
||||
),
|
||||
output_token_ids=output_token_ids,
|
||||
spec_token_ids=[[] for _ in range(batch_size)],
|
||||
frequency_penalties=_create_penalty_tensor(batch_size, 0.0, device),
|
||||
presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),
|
||||
repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),
|
||||
no_penalties=True,
|
||||
allowed_token_ids_mask=None,
|
||||
bad_words_token_ids={},
|
||||
logitsprocs=LogitsProcessors(),
|
||||
)
|
||||
return fake_sampling_metadata
|
||||
|
||||
|
||||
def _create_weighted_output_token_list(
|
||||
batch_size: int, vocab_size: int
|
||||
) -> tuple[list[list[int]], list[list[int]]]:
|
||||
"""
|
||||
Creates an output token list where each token occurs a distinct
|
||||
number of times.
|
||||
|
||||
For each batch, a random subset of token IDs is selected from the
|
||||
vocabulary. The selected tokens are then added to the output token
|
||||
list, each with a different frequency.
|
||||
|
||||
Returns:
|
||||
tuple[list[list[int]], list[list[int]]]:
|
||||
- The first element is the output token list, where each sublist
|
||||
corresponds to a batch and contains tokens with weighted
|
||||
frequencies.
|
||||
- The second element is a list of distinct token IDs for each
|
||||
batch, ordered by their frequency in the corresponding output
|
||||
list.
|
||||
"""
|
||||
output_token_ids: list[list[int]] = []
|
||||
sorted_token_ids_in_output: list[list[int]] = []
|
||||
for _ in range(batch_size):
|
||||
distinct_token_ids = np.random.choice(
|
||||
vocab_size, size=np.random.randint(1, 10), replace=False
|
||||
).tolist()
|
||||
sorted_token_ids_in_output.append(distinct_token_ids)
|
||||
output_token_ids_for_batch = []
|
||||
for index, token_id in enumerate(distinct_token_ids):
|
||||
output_token_ids_for_batch.extend([token_id for _ in range(index + 1)])
|
||||
output_token_ids.append(output_token_ids_for_batch)
|
||||
return output_token_ids, sorted_token_ids_in_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 32])
|
||||
@pytest.mark.parametrize("presence_penalty", [-2.0, 2.0])
|
||||
def test_sampler_presence_penalty(
|
||||
device: str, batch_size: int, presence_penalty: float
|
||||
):
|
||||
"""
|
||||
Test to verify that if presence penalty is enabled then tokens
|
||||
are penalized as per their presence in the existing output.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
# Create fake logits where each token is assigned the same
|
||||
# logit value.
|
||||
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)
|
||||
)
|
||||
output_token_ids = sampling_metadata.output_token_ids
|
||||
sampling_metadata.presence_penalties = _create_penalty_tensor(
|
||||
batch_size, presence_penalty, torch.device(device)
|
||||
)
|
||||
sampling_metadata.no_penalties = False
|
||||
sampler = Sampler()
|
||||
logits = sampler.apply_penalties(
|
||||
fake_logits, sampling_metadata, sampling_metadata.output_token_ids
|
||||
)
|
||||
logits = logits.cpu()
|
||||
for batch_idx in range(batch_size):
|
||||
# Since all tokens initially have the same logits, the non-penalized
|
||||
# token ID will be the one with the highest logit value, while the
|
||||
# penalized token ID will be the one with the lowest logit value.
|
||||
non_penalized_token_id = logits[batch_idx].argmax().item()
|
||||
penalized_token_id = logits[batch_idx].argmin().item()
|
||||
if presence_penalty > 0:
|
||||
# If `presence_penalty` is set to a value greater than 0, it
|
||||
# indicates a preference for new tokens over those already
|
||||
# present in the output.
|
||||
# Verify that the penalized token ID exists in the output, while the
|
||||
# non-penalized token ID does not.
|
||||
assert penalized_token_id in output_token_ids[batch_idx]
|
||||
assert non_penalized_token_id not in output_token_ids[batch_idx]
|
||||
elif presence_penalty < 0:
|
||||
# If `presence_penalty` is set to a value less than 0, it indicates
|
||||
# a preference for existing tokens over new ones. Verify that the
|
||||
# non-penalized token ID exists in the output, while the penalized
|
||||
# token ID does not.
|
||||
assert non_penalized_token_id in output_token_ids[batch_idx]
|
||||
assert penalized_token_id not in output_token_ids[batch_idx]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 32])
|
||||
@pytest.mark.parametrize("frequency_penalty", [-2.0, 2.0])
|
||||
def test_sampler_frequency_penalty(
|
||||
device: str, batch_size: int, frequency_penalty: float
|
||||
):
|
||||
"""
|
||||
Test to verify that if frequency penalty is enabled then tokens are
|
||||
penalized as per their frequency of occurrence.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
# Create fake logits where each token is assigned the same
|
||||
# logit value.
|
||||
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)
|
||||
)
|
||||
sampling_metadata.frequency_penalties = _create_penalty_tensor(
|
||||
batch_size, frequency_penalty, torch.device(device)
|
||||
)
|
||||
output_token_ids, sorted_token_ids_in_output = _create_weighted_output_token_list(
|
||||
batch_size,
|
||||
VOCAB_SIZE,
|
||||
)
|
||||
sampling_metadata.output_token_ids = output_token_ids
|
||||
sampling_metadata.no_penalties = False
|
||||
sampler = Sampler()
|
||||
logits = sampler.apply_penalties(
|
||||
fake_logits, sampling_metadata, sampling_metadata.output_token_ids
|
||||
)
|
||||
logits = logits.cpu()
|
||||
for batch_idx in range(batch_size):
|
||||
non_penalized_token_id = logits[batch_idx].argmax().item()
|
||||
penalized_token_id = logits[batch_idx].argmin().item()
|
||||
distinct_sorted_token_ids_in_output = sorted_token_ids_in_output[batch_idx]
|
||||
most_frequent_token_id = distinct_sorted_token_ids_in_output[
|
||||
len(distinct_sorted_token_ids_in_output) - 1
|
||||
]
|
||||
if frequency_penalty > 0:
|
||||
# If `frequency_penalty` is set to > 0, it indicates
|
||||
# a preference for new tokens over existing ones. Verify that the
|
||||
# non-penalized token ID is not present in the output, while the
|
||||
# most penalized token is the one that occurs most frequently in
|
||||
# the output.
|
||||
assert non_penalized_token_id not in distinct_sorted_token_ids_in_output
|
||||
assert penalized_token_id == most_frequent_token_id
|
||||
elif frequency_penalty < 0:
|
||||
# If `frequency_penalty` is set to < 0, it indicates
|
||||
# a preference for existing tokens over new ones. Verify that the
|
||||
# non-penalized token ID is the one that occurs most frequently
|
||||
# in the output, while the penalized token ID is one that has not
|
||||
# yet appeared.
|
||||
assert non_penalized_token_id == most_frequent_token_id
|
||||
assert penalized_token_id not in distinct_sorted_token_ids_in_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 32])
|
||||
@pytest.mark.parametrize("repetition_penalty", [0.1, 1.9])
|
||||
def test_sampler_repetition_penalty(
|
||||
device: str, batch_size: int, repetition_penalty: float
|
||||
):
|
||||
"""
|
||||
Test to verify that when the repetition penalty is enabled, tokens
|
||||
are penalized based on their presence in the prompt or the existing
|
||||
output.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
# Create fake logits where each token is assigned the same
|
||||
# logit value.
|
||||
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)
|
||||
)
|
||||
sampling_metadata.repetition_penalties = _create_penalty_tensor(
|
||||
batch_size, repetition_penalty, torch.device(device)
|
||||
)
|
||||
sampling_metadata.no_penalties = False
|
||||
sampler = Sampler()
|
||||
logits = sampler.apply_penalties(
|
||||
fake_logits, sampling_metadata, sampling_metadata.output_token_ids
|
||||
)
|
||||
logits = logits.cpu()
|
||||
for batch_idx in range(batch_size):
|
||||
non_penalized_token_id = logits[batch_idx].argmax().item()
|
||||
penalized_token_id = logits[batch_idx].argmin().item()
|
||||
prompt_tokens = sampling_metadata.prompt_token_ids[batch_idx][:].tolist()
|
||||
output_tokens = sampling_metadata.output_token_ids[batch_idx]
|
||||
if repetition_penalty > 1.0:
|
||||
# If `repetition_penalty` > 1.0, verify that the non-penalized
|
||||
# token ID has not been seen before, while the penalized token ID
|
||||
# exists either in the prompt or the output.
|
||||
assert (
|
||||
non_penalized_token_id not in prompt_tokens
|
||||
and non_penalized_token_id not in output_tokens
|
||||
)
|
||||
assert (
|
||||
penalized_token_id in prompt_tokens
|
||||
or penalized_token_id in output_tokens
|
||||
)
|
||||
elif repetition_penalty < 1.0:
|
||||
# If `repetition_penalty` < 1.0, verify that the penalized
|
||||
# token ID has not been seen before, while the non-penalized
|
||||
# token ID exists either in the prompt or the output.
|
||||
assert (
|
||||
penalized_token_id not in prompt_tokens
|
||||
and penalized_token_id not in output_tokens
|
||||
)
|
||||
assert (
|
||||
non_penalized_token_id in prompt_tokens
|
||||
or non_penalized_token_id in output_tokens
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 32])
|
||||
@pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])
|
||||
def test_sampler_allowed_token_ids(
|
||||
device: str, batch_size: int, num_allowed_token_ids: int
|
||||
):
|
||||
"""
|
||||
Test to verify that when the repetition penalty is enabled, tokens
|
||||
are penalized based on their presence in the prompt or the existing
|
||||
output.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
# Create fake logits where each token is assigned the same
|
||||
# logit value.
|
||||
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)
|
||||
)
|
||||
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.allowed_token_ids_mask = mask
|
||||
sampler = Sampler()
|
||||
logits = sampler.apply_logits_processors(
|
||||
fake_logits, sampling_metadata, predict_bonus_token=False
|
||||
)
|
||||
logits = logits.cpu()
|
||||
for batch_idx in range(batch_size):
|
||||
logits_for_req = logits[batch_idx]
|
||||
if batch_idx % 2 == 1:
|
||||
assert torch.all(logits_for_req != -float("inf"))
|
||||
continue
|
||||
for token_id in range(VOCAB_SIZE):
|
||||
start = min(batch_idx, VOCAB_SIZE - 1)
|
||||
end = min(batch_idx + num_allowed_token_ids, VOCAB_SIZE - 1)
|
||||
if token_id >= start and token_id < end:
|
||||
assert logits_for_req[token_id] == -float("inf"), (
|
||||
f"{batch_idx}, {token_id}"
|
||||
)
|
||||
else:
|
||||
assert logits_for_req[token_id] != -float("inf")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 32])
|
||||
@pytest.mark.parametrize("bad_words_lengths", [(1,), (1, 3), (2, 2)])
|
||||
def test_sampler_bad_words(
|
||||
device: str, batch_size: int, bad_words_lengths: tuple[int, ...]
|
||||
):
|
||||
"""
|
||||
Test to verify that when the bad words restriction is present, tokens
|
||||
are penalized based on their match with the bad words.
|
||||
"""
|
||||
torch.set_default_device(device)
|
||||
# Create fake logits where each token is assigned the same
|
||||
# logit value.
|
||||
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)
|
||||
)
|
||||
sampling_metadata.bad_words_token_ids = _create_bad_words_token_ids(
|
||||
batch_size, VOCAB_SIZE, bad_words_lengths
|
||||
)
|
||||
bad_words_last_tokens = _update_output_token_ids_for_bad_words(
|
||||
sampling_metadata, VOCAB_SIZE
|
||||
)
|
||||
sampler = Sampler()
|
||||
logits = sampler.apply_logits_processors(
|
||||
fake_logits, sampling_metadata, predict_bonus_token=False
|
||||
)
|
||||
logits = logits.cpu()
|
||||
for batch_idx in range(batch_size):
|
||||
logits_for_req = logits[batch_idx]
|
||||
for token_id in range(VOCAB_SIZE):
|
||||
if (
|
||||
batch_idx in bad_words_last_tokens
|
||||
and token_id in bad_words_last_tokens[batch_idx]
|
||||
):
|
||||
assert logits_for_req[token_id] == -float("inf")
|
||||
else:
|
||||
assert logits_for_req[token_id] != -float("inf")
|
||||
@@ -0,0 +1,184 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
PROMPT = "Hello my name is Robert and I"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm() -> LLM:
|
||||
return LLM(MODEL, enforce_eager=True)
|
||||
|
||||
|
||||
def test_n_gt_1(llm):
|
||||
"""ParallelSampling is supported."""
|
||||
|
||||
params = SamplingParams(n=3)
|
||||
outputs = llm.generate(PROMPT, params)
|
||||
assert len(outputs[0].outputs) == 3
|
||||
|
||||
|
||||
def test_penalties(llm):
|
||||
"""Check that we do not get errors if applied."""
|
||||
|
||||
params = SamplingParams(
|
||||
temperature=1.2,
|
||||
presence_penalty=1.2,
|
||||
frequency_penalty=1.2,
|
||||
repetition_penalty=1.2,
|
||||
min_p=0.5,
|
||||
top_p=0.5,
|
||||
top_k=3,
|
||||
)
|
||||
_ = llm.generate(PROMPT, params)
|
||||
|
||||
|
||||
def test_stop(llm):
|
||||
"""Check that we respect the stop words."""
|
||||
|
||||
output = llm.generate(PROMPT, SamplingParams(temperature=0))
|
||||
split_text = output[0].outputs[0].text.split()
|
||||
|
||||
STOP_IDX = 5
|
||||
params = SamplingParams(temperature=0, stop=split_text[STOP_IDX])
|
||||
output = llm.generate(PROMPT, params)
|
||||
new_split_text = output[0].outputs[0].text.split()
|
||||
|
||||
# Output should not contain the stop word.
|
||||
assert len(new_split_text) == STOP_IDX
|
||||
|
||||
params = SamplingParams(
|
||||
temperature=0, stop=split_text[STOP_IDX], include_stop_str_in_output=True
|
||||
)
|
||||
output = llm.generate(PROMPT, params)
|
||||
new_split_text = output[0].outputs[0].text.split()
|
||||
|
||||
# Output should contain the stop word.
|
||||
assert len(new_split_text) == STOP_IDX + 1
|
||||
|
||||
|
||||
def test_stop_token_ids(llm):
|
||||
"""Check that we respect the stop token ids."""
|
||||
|
||||
output = llm.generate(PROMPT, SamplingParams(temperature=0))
|
||||
|
||||
stop_token_id_0 = output[0].outputs[0].token_ids[5]
|
||||
stop_token_id_1 = output[0].outputs[0].token_ids[6]
|
||||
|
||||
stop_token_ids = [stop_token_id_1, stop_token_id_0]
|
||||
params = SamplingParams(temperature=0, stop_token_ids=stop_token_ids)
|
||||
output = llm.generate(PROMPT, params)
|
||||
assert output[0].outputs[0].token_ids[-1] == stop_token_id_0
|
||||
|
||||
stop_token_ids = [stop_token_id_0, stop_token_id_1]
|
||||
params = SamplingParams(temperature=0, stop_token_ids=stop_token_ids)
|
||||
output = llm.generate(PROMPT, params)
|
||||
assert output[0].outputs[0].token_ids[-1] == stop_token_id_0
|
||||
|
||||
|
||||
def test_detokenize_false(llm):
|
||||
"""Check that detokenize=False option works."""
|
||||
|
||||
output = llm.generate(PROMPT, SamplingParams(detokenize=False))
|
||||
assert len(output[0].outputs[0].token_ids) > 0
|
||||
assert len(output[0].outputs[0].text) == 0
|
||||
|
||||
output = llm.generate(
|
||||
PROMPT, SamplingParams(detokenize=False, logprobs=3, prompt_logprobs=3)
|
||||
)
|
||||
assert len(output[0].outputs[0].token_ids) > 0
|
||||
assert len(output[0].outputs[0].text) == 0
|
||||
|
||||
prompt_logprobs = output[0].prompt_logprobs
|
||||
sampled_logprobs = output[0].outputs[0].logprobs
|
||||
assert len(prompt_logprobs) > 1
|
||||
assert len(sampled_logprobs) > 1
|
||||
for all_logprobs in (prompt_logprobs[1:], sampled_logprobs):
|
||||
for logprobs in all_logprobs:
|
||||
assert 3 <= len(logprobs) <= 4
|
||||
assert all(lp.decoded_token is None for lp in logprobs.values())
|
||||
|
||||
|
||||
def test_bad_words(llm):
|
||||
"""Check that we respect bad words."""
|
||||
|
||||
tokenizer = llm.get_tokenizer()
|
||||
|
||||
def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool:
|
||||
"""Check if word appears in BOTH text and token sequence."""
|
||||
if bad_word not in text:
|
||||
return False
|
||||
|
||||
for add_prefix_space in [False, True]:
|
||||
prefix = " " if add_prefix_space else ""
|
||||
bad_words_token = tokenizer.encode(
|
||||
prefix + bad_word.lstrip(), add_special_tokens=False
|
||||
)
|
||||
if not bad_words_token:
|
||||
continue
|
||||
for i in range(len(tokens) - len(bad_words_token) + 1):
|
||||
if tokens[i : i + len(bad_words_token)] == bad_words_token:
|
||||
return True
|
||||
return False
|
||||
|
||||
output = llm.generate(PROMPT, SamplingParams(temperature=0))
|
||||
split_text = output[0].outputs[0].text.split()
|
||||
|
||||
bad_words_1 = " ".join(split_text[:2])
|
||||
params = SamplingParams(temperature=0, bad_words=[bad_words_1])
|
||||
output = llm.generate(PROMPT, params)
|
||||
new_text = output[0].outputs[0].text
|
||||
new_tokens = output[0].outputs[0].token_ids
|
||||
assert not contains_bad_word(new_text, new_tokens, bad_words_1)
|
||||
|
||||
bad_words_2 = new_text.split()[-1]
|
||||
params = SamplingParams(temperature=0, bad_words=[bad_words_1, bad_words_2])
|
||||
output = llm.generate(PROMPT, params)
|
||||
new_text = output[0].outputs[0].text
|
||||
new_tokens = output[0].outputs[0].token_ids
|
||||
assert not contains_bad_word(new_text, new_tokens, bad_words_1)
|
||||
assert not contains_bad_word(new_text, new_tokens, bad_words_2)
|
||||
|
||||
|
||||
def test_allowed_token_ids(llm):
|
||||
"""Check that we can use allowed_token_ids."""
|
||||
|
||||
TOKEN_ID = 10
|
||||
allowed_token_ids = [TOKEN_ID]
|
||||
output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids))
|
||||
assert output[0].outputs[0].token_ids[-1] == TOKEN_ID
|
||||
|
||||
# Each single-token allowlist must force that token (kernel used to drop some).
|
||||
for token_id in (1, 5, 100, 500, 2518, 9834, 31999):
|
||||
output = llm.generate(
|
||||
PROMPT,
|
||||
SamplingParams(temperature=0, max_tokens=1, allowed_token_ids=[token_id]),
|
||||
)
|
||||
assert output[0].outputs[0].token_ids[-1] == token_id
|
||||
|
||||
# Reject empty allowed_token_ids.
|
||||
with pytest.raises(ValueError):
|
||||
_ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[]))
|
||||
|
||||
# Reject negative token id.
|
||||
with pytest.raises(ValueError):
|
||||
_ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[-1]))
|
||||
|
||||
# Reject out of vocabulary.
|
||||
with pytest.raises(ValueError):
|
||||
_ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[10000000]))
|
||||
|
||||
|
||||
def test_seed(llm):
|
||||
"""Check that seed impacts randomness."""
|
||||
|
||||
out_1 = llm.generate(PROMPT, SamplingParams(seed=42))
|
||||
out_2 = llm.generate(PROMPT, SamplingParams(seed=42))
|
||||
out_3 = llm.generate(PROMPT, SamplingParams(seed=43))
|
||||
|
||||
assert out_1[0].outputs[0].text == out_2[0].outputs[0].text
|
||||
assert out_1[0].outputs[0].text != out_3[0].outputs[0].text
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,258 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterator
|
||||
from enum import Enum
|
||||
from typing import NamedTuple
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
from vllm import CompletionOutput
|
||||
from vllm.utils.torch_utils import make_tensor_with_pad
|
||||
from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
|
||||
|
||||
class BatchLogprobsComposition(Enum):
|
||||
"""Types of logprobs configs to include in test batch"""
|
||||
|
||||
NONE = 0
|
||||
SAMPLE = 1
|
||||
PROMPT = 2
|
||||
SAMPLE_PROMPT = 3
|
||||
|
||||
|
||||
BatchLogprobsSpecType = list[tuple[int | None, int | None]]
|
||||
|
||||
|
||||
def get_test_batch(
|
||||
batch_logprobs_composition: BatchLogprobsComposition,
|
||||
) -> BatchLogprobsSpecType:
|
||||
"""Generate logprobs configs for a batch of requests
|
||||
|
||||
A given request's logprobs configuration is (1) num_sample_logprobs and (2)
|
||||
num_prompt_logprobs. The batch logprobs configuration is the list of request
|
||||
logprobs configs.
|
||||
|
||||
batch_logprobs_composition == NONE yields a batch with no sample or prompt
|
||||
logprobs
|
||||
|
||||
batch_logprobs_composition == SAMPLE yields a batch with some requests
|
||||
configured for sample logprobs only, and others configured for no logprobs
|
||||
|
||||
batch_logprobs_composition == PROMPT yields a batch with some requests
|
||||
configured for prompt logprobs only, and others configured for no logprobs
|
||||
|
||||
batch_logprobs_composition == SAMPLE_PROMPT yields a batch with some
|
||||
requests configured for sample logprobs and prompt logprobs, some configured
|
||||
for only sample logprobs or only prompt logprobs, and some configured for
|
||||
no logprobs
|
||||
|
||||
Args:
|
||||
batch_logprobs_composition: types of logprobs configs to include in batch
|
||||
|
||||
Returns:
|
||||
|
||||
list of (Optional[num_sample_logprobs], Optional[num_prompt_logprobs])
|
||||
tuples
|
||||
"""
|
||||
if batch_logprobs_composition == BatchLogprobsComposition.NONE:
|
||||
# No requests with sample or prompt logprobs
|
||||
return [(None, None)]
|
||||
elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE:
|
||||
# Requests requiring sample logprobs or no logprobs
|
||||
return [
|
||||
(None, None),
|
||||
(0, None),
|
||||
(5, None),
|
||||
(3, None),
|
||||
]
|
||||
elif batch_logprobs_composition == BatchLogprobsComposition.PROMPT:
|
||||
# Requests requiring prompt logprobs or no logprobs
|
||||
return [
|
||||
(None, None),
|
||||
(None, 0),
|
||||
(None, 6),
|
||||
(None, 5),
|
||||
]
|
||||
elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE_PROMPT:
|
||||
# Requests requiring either no logprobs, just
|
||||
# sample logprobs, just prompt logprobs, or
|
||||
# both sample and prompt logprobs
|
||||
return [
|
||||
(None, None),
|
||||
(0, None),
|
||||
(5, None),
|
||||
(3, None),
|
||||
(0, 3),
|
||||
(6, 0),
|
||||
(6, 3),
|
||||
(None, 6),
|
||||
(None, 5),
|
||||
(None, 0),
|
||||
]
|
||||
else:
|
||||
raise ValueError("Invalid logprobs batch configuration for test.")
|
||||
|
||||
|
||||
def assert_incr_detok_str_matches_non_incr_detok_str(
|
||||
incremental_detokenization_str: str,
|
||||
non_incremental_detokenization_str: str,
|
||||
msg: str,
|
||||
) -> None:
|
||||
"""Compare incrementally detok. text to non-incrementally detok. text
|
||||
|
||||
Fail if the strings mismatch after non-alphanumeric characters are stripped
|
||||
out.
|
||||
|
||||
Rationale: incremental detokenization in the text generation process allows
|
||||
the tokenizer to adjust the next token text output based on the token's
|
||||
context in the string. However, logprobs detokenization detokenizes each
|
||||
token individually, and the resultant strings may include some
|
||||
non-alphanumeric placeholder characters where there could be i.e.
|
||||
whitespace. So, this function compares only the alphanumeric text
|
||||
between two strings and fails if there is a mismatch, which helps
|
||||
with validating logprobs detokenization.
|
||||
|
||||
Args:
|
||||
incremental_detokenization_str: incrementally-detokenized generated text
|
||||
non_incremental_detokenization_str: non-incrementally-detokenized logprob
|
||||
tokens
|
||||
msg: error message if `assert` fails
|
||||
"""
|
||||
rgx = r"[^a-zA-Z0-9]+"
|
||||
assert re.sub(rgx, "", incremental_detokenization_str) == re.sub(
|
||||
rgx, "", non_incremental_detokenization_str
|
||||
), msg
|
||||
|
||||
|
||||
def compute_correct_cumulative_logprob(completion_output: CompletionOutput) -> float:
|
||||
"""Compute known-good value for evaluating cumulative logprob
|
||||
|
||||
Args:
|
||||
completion_output: completion output from engine
|
||||
|
||||
Returns:
|
||||
Known-good cumulative logprob value
|
||||
"""
|
||||
token_ids = completion_output.token_ids
|
||||
logprobs = completion_output.logprobs
|
||||
assert logprobs is not None
|
||||
return sum([lp[tok_id].logprob for tok_id, lp in zip(token_ids, logprobs)])
|
||||
|
||||
|
||||
def create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor:
|
||||
fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float)
|
||||
return fake_logits
|
||||
|
||||
|
||||
def create_penalty_tensor(
|
||||
batch_size: int, penalty_value: float, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
return torch.full(
|
||||
(batch_size,), fill_value=penalty_value, dtype=torch.float, device=device
|
||||
)
|
||||
|
||||
|
||||
def create_prompt_tokens_tensor(
|
||||
prompt_token_ids: list[list[int]],
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
return make_tensor_with_pad(
|
||||
prompt_token_ids,
|
||||
pad=vocab_size,
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
pin_memory=False,
|
||||
)
|
||||
|
||||
|
||||
class LogitsprocsTestFakes(NamedTuple):
|
||||
"""Wraps fake data structures to support testing"""
|
||||
|
||||
logits: torch.Tensor
|
||||
sampling_metadata: SamplingMetadata
|
||||
|
||||
def get_logitsprocs_by_cls(
|
||||
self,
|
||||
cls: type[LogitsProcessor],
|
||||
) -> Iterator[LogitsProcessor]:
|
||||
"""Yield logits processors of a specific class.
|
||||
|
||||
Args:
|
||||
cls: :class:`LogitsProcessor` subclass
|
||||
|
||||
Returns:
|
||||
Iterator over logits processors
|
||||
"""
|
||||
return (
|
||||
lp for lp in self.sampling_metadata.logitsprocs.all if isinstance(lp, cls)
|
||||
)
|
||||
|
||||
def get_logitsprocs(self) -> Iterator[LogitsProcessor]:
|
||||
"""Iterator over all logits processors."""
|
||||
return self.sampling_metadata.logitsprocs.all
|
||||
|
||||
|
||||
def fake_update_logitsprocs_state(
|
||||
test_fakes: LogitsprocsTestFakes,
|
||||
batch_update: BatchUpdate | None,
|
||||
) -> None:
|
||||
"""Imitate logits processors persistent batch state update
|
||||
in engine core"""
|
||||
for logitproc in test_fakes.get_logitsprocs():
|
||||
logitproc.update_state(batch_update)
|
||||
holder = test_fakes.sampling_metadata.thinking_budget_state_holder
|
||||
if holder is not None:
|
||||
holder.sync_batch(batch_update)
|
||||
|
||||
|
||||
def fake_apply_logitsprocs(
|
||||
test_fakes: LogitsprocsTestFakes,
|
||||
slice_indices: list[int],
|
||||
slot_output_token_ids: list[list[int]] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Imitate application of logits processors in engine core.
|
||||
|
||||
When ``thinking_budget_state_holder`` has tracked requests, this mirrors
|
||||
:meth:`Sampler.apply_logits_processors` by refreshing per-slot
|
||||
``output_token_ids`` (if ``slot_output_token_ids`` is provided), then
|
||||
``update_state`` + ``apply_to_logits`` on the holder after built-in logits
|
||||
processors.
|
||||
"""
|
||||
logits = test_fakes.logits[torch.tensor(slice_indices, dtype=torch.long)].clone()
|
||||
for processor in test_fakes.get_logitsprocs():
|
||||
logits = processor.apply(logits)
|
||||
|
||||
md = test_fakes.sampling_metadata
|
||||
holder = md.thinking_budget_state_holder
|
||||
if holder is not None and holder.has_tracked_requests():
|
||||
if slot_output_token_ids is not None:
|
||||
for i, toks in enumerate(slot_output_token_ids):
|
||||
if i < len(md.output_token_ids):
|
||||
md.output_token_ids[i] = list(toks)
|
||||
holder.update_state(md.output_token_ids, md.spec_token_ids, None)
|
||||
logits = holder.apply_to_logits(logits, False, md.spec_token_ids)
|
||||
return logits
|
||||
|
||||
|
||||
def create_allowed_token_ids(
|
||||
batch_size: int,
|
||||
vocab_size: int,
|
||||
num_allowed_token_ids: int,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor | None:
|
||||
mask: torch.Tensor | None = None
|
||||
for i in range(batch_size):
|
||||
if i % 2 == 1:
|
||||
continue
|
||||
if mask is None:
|
||||
mask = torch.zeros(
|
||||
(batch_size, vocab_size), dtype=torch.bool, device=device
|
||||
)
|
||||
start = min(i, vocab_size - 1)
|
||||
end = min(i + num_allowed_token_ids, vocab_size - 1)
|
||||
mask[i, start:end] = True
|
||||
return mask
|
||||
Reference in New Issue
Block a user