init
This commit is contained in:
268
vllm_vacc/vllm/model_executor/sampling_metadata.py
Normal file
268
vllm_vacc/vllm/model_executor/sampling_metadata.py
Normal file
@@ -0,0 +1,268 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from array import array
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.sampling_params import SamplingParams, SamplingType
|
||||
from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE,
|
||||
SequenceGroupMetadata)
|
||||
from vllm.utils import (PyObjectCache, async_tensor_h2d,
|
||||
is_pin_memory_available, make_tensor_with_pad)
|
||||
|
||||
|
||||
from vllm.model_executor.sampling_metadata import SamplingTensors, SamplingMetadataCache, _prepare_seq_groups, SamplingMetadata, _SAMPLING_EPS
|
||||
|
||||
@staticmethod
|
||||
def SamplingMetadata_prepare(
|
||||
seq_group_metadata_list: List[SequenceGroupMetadata],
|
||||
seq_lens: List[int],
|
||||
query_lens: List[int],
|
||||
device: str,
|
||||
pin_memory: bool,
|
||||
generators: Optional[Dict[str, torch.Generator]] = None,
|
||||
cache: Optional[SamplingMetadataCache] = None,
|
||||
) -> "SamplingMetadata":
|
||||
(
|
||||
seq_groups,
|
||||
selected_token_indices,
|
||||
categorized_sample_indices,
|
||||
num_prompts,
|
||||
) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,
|
||||
device, generators, cache)
|
||||
selected_token_indices = async_tensor_h2d(
|
||||
selected_token_indices,
|
||||
dtype=torch.int32, #use int32 instead of long
|
||||
target_device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
categorized_sample_indices = {
|
||||
t:
|
||||
async_tensor_h2d(
|
||||
seq_ids,
|
||||
dtype=torch.int,
|
||||
target_device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for t, seq_ids in categorized_sample_indices.items()
|
||||
}
|
||||
|
||||
sampling_metadata = SamplingMetadata(
|
||||
seq_groups=seq_groups,
|
||||
selected_token_indices=selected_token_indices,
|
||||
categorized_sample_indices=categorized_sample_indices,
|
||||
num_prompts=num_prompts,
|
||||
)
|
||||
return sampling_metadata
|
||||
|
||||
@classmethod
|
||||
def SamplingTensors_from_lists(
|
||||
cls,
|
||||
temperatures: List[float],
|
||||
top_ps: List[float],
|
||||
top_ks: List[int],
|
||||
min_ps: List[float],
|
||||
presence_penalties: List[float],
|
||||
frequency_penalties: List[float],
|
||||
repetition_penalties: List[float],
|
||||
prompt_tokens: List[array],
|
||||
output_tokens: List[array],
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> "SamplingTensors":
|
||||
# Note that the performance will be very bad without
|
||||
# pinned memory.
|
||||
pin_memory = is_pin_memory_available()
|
||||
|
||||
do_penalties = prompt_tokens or output_tokens
|
||||
|
||||
if do_penalties:
|
||||
prompt_t = make_tensor_with_pad(
|
||||
prompt_tokens,
|
||||
vocab_size,
|
||||
device="cpu",
|
||||
dtype=torch.int64,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
output_t = make_tensor_with_pad(
|
||||
output_tokens,
|
||||
vocab_size,
|
||||
device="cpu",
|
||||
dtype=torch.int64,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
else:
|
||||
empty_tensor = torch.empty(0, device=device, dtype=torch.long)
|
||||
prompt_t = empty_tensor
|
||||
output_t = empty_tensor
|
||||
|
||||
temperatures_t = torch.tensor(
|
||||
temperatures,
|
||||
device="cpu",
|
||||
dtype=torch.float32,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
top_ps_t = torch.tensor(
|
||||
top_ps,
|
||||
device="cpu",
|
||||
dtype=torch.float32,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
min_ps_t = torch.tensor(
|
||||
min_ps,
|
||||
device="cpu",
|
||||
dtype=dtype,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
presence_penalties_t = torch.tensor(
|
||||
presence_penalties,
|
||||
device="cpu",
|
||||
dtype=dtype,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
frequency_penalties_t = torch.tensor(
|
||||
frequency_penalties,
|
||||
device="cpu",
|
||||
dtype=dtype,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
repetition_penalties_t = torch.tensor(
|
||||
repetition_penalties,
|
||||
device="cpu",
|
||||
dtype=dtype,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
top_ks_t = torch.tensor(
|
||||
top_ks,
|
||||
device="cpu",
|
||||
dtype=torch.int,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
# Because the memory is pinned, we can do non-blocking
|
||||
return cls(
|
||||
temperatures=temperatures_t,
|
||||
top_ps=top_ps_t,
|
||||
top_ks=top_ks_t,
|
||||
min_ps=min_ps_t,
|
||||
presence_penalties=presence_penalties_t,
|
||||
frequency_penalties=frequency_penalties_t,
|
||||
repetition_penalties=repetition_penalties_t,
|
||||
prompt_tokens=prompt_t,
|
||||
output_tokens=output_t,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def SamplingMetadata_from_sampling_metadata(
|
||||
cls,
|
||||
sampling_metadata: "SamplingMetadata",
|
||||
vocab_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> Tuple["SamplingTensors", bool, bool, bool]:
|
||||
prompt_tokens: List[array] = []
|
||||
output_tokens: List[array] = []
|
||||
top_ks: List[int] = []
|
||||
temperatures: List[float] = []
|
||||
top_ps: List[float] = []
|
||||
min_ps: List[float] = []
|
||||
presence_penalties: List[float] = []
|
||||
frequency_penalties: List[float] = []
|
||||
repetition_penalties: List[float] = []
|
||||
do_penalties = False
|
||||
do_top_p_top_k = False
|
||||
do_min_p = False
|
||||
|
||||
assert sampling_metadata.seq_groups is not None
|
||||
for seq_group in sampling_metadata.seq_groups:
|
||||
seq_ids = seq_group.seq_ids
|
||||
sampling_params = seq_group.sampling_params
|
||||
temperature = sampling_params.temperature
|
||||
p = sampling_params.presence_penalty
|
||||
f = sampling_params.frequency_penalty
|
||||
r = sampling_params.repetition_penalty
|
||||
top_p = sampling_params.top_p
|
||||
min_p = sampling_params.min_p
|
||||
|
||||
# k should not be greater than the vocab size.
|
||||
top_k = min(sampling_params.top_k, vocab_size)
|
||||
# top_k = vocab_size if top_k == -1 else top_k
|
||||
# FIXME: fix top_k to avoid odsp bug currently
|
||||
top_k = 40
|
||||
if temperature < _SAMPLING_EPS:
|
||||
# NOTE: Zero temperature means deterministic sampling
|
||||
# (i.e., greedy sampling or beam search).
|
||||
# Set the temperature to 1 to avoid division by zero.
|
||||
temperature = 1.0
|
||||
if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
|
||||
or top_k != vocab_size):
|
||||
do_top_p_top_k = True
|
||||
if not do_min_p and min_p > _SAMPLING_EPS:
|
||||
do_min_p = True
|
||||
if not do_penalties and (abs(p) >= _SAMPLING_EPS
|
||||
or abs(f) >= _SAMPLING_EPS
|
||||
or abs(r - 1.0) >= _SAMPLING_EPS):
|
||||
do_penalties = True
|
||||
|
||||
is_prompt = seq_group.is_prompt
|
||||
if is_prompt and sampling_params.prompt_logprobs is not None:
|
||||
# For tokens in the prompt that we only need to get
|
||||
# their logprobs
|
||||
query_len = seq_group.query_len
|
||||
assert query_len is not None
|
||||
prefill_len = len(seq_group.prompt_logprob_indices)
|
||||
temperatures += [temperature] * prefill_len
|
||||
top_ps += [top_p] * prefill_len
|
||||
top_ks += [top_k] * prefill_len
|
||||
min_ps += [min_p] * prefill_len
|
||||
presence_penalties += [0] * prefill_len
|
||||
frequency_penalties += [0] * prefill_len
|
||||
repetition_penalties += [1] * prefill_len
|
||||
|
||||
if seq_group.do_sample:
|
||||
sample_lens = len(seq_group.sample_indices)
|
||||
assert sample_lens >= len(seq_ids)
|
||||
temperatures += [temperature] * sample_lens
|
||||
top_ps += [top_p] * sample_lens
|
||||
top_ks += [top_k] * sample_lens
|
||||
min_ps += [min_p] * sample_lens
|
||||
presence_penalties += [p] * sample_lens
|
||||
frequency_penalties += [f] * sample_lens
|
||||
repetition_penalties += [r] * sample_lens
|
||||
|
||||
if do_penalties:
|
||||
for seq_group in sampling_metadata.seq_groups:
|
||||
seq_ids = seq_group.seq_ids
|
||||
sampling_params = seq_group.sampling_params
|
||||
if (seq_group.is_prompt
|
||||
and sampling_params.prompt_logprobs is not None):
|
||||
prefill_len = len(seq_group.prompt_logprob_indices)
|
||||
prompt_tokens.extend(
|
||||
array(VLLM_TOKEN_ID_ARRAY_TYPE)
|
||||
for _ in range(prefill_len))
|
||||
output_tokens.extend(
|
||||
array(VLLM_TOKEN_ID_ARRAY_TYPE)
|
||||
for _ in range(prefill_len))
|
||||
if seq_group.do_sample:
|
||||
for seq_id in seq_ids:
|
||||
seq_data = seq_group.seq_data[seq_id]
|
||||
prompt_tokens.append(seq_data.prompt_token_ids_array)
|
||||
output_tokens.append(seq_data.output_token_ids_array)
|
||||
|
||||
sampling_tensors = SamplingTensors.from_lists(
|
||||
temperatures,
|
||||
top_ps,
|
||||
top_ks,
|
||||
min_ps,
|
||||
presence_penalties,
|
||||
frequency_penalties,
|
||||
repetition_penalties,
|
||||
prompt_tokens,
|
||||
output_tokens,
|
||||
vocab_size,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)
|
||||
Reference in New Issue
Block a user