mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-01-21 22:14:41 +08:00
127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch_xla.core.xla_model as xm
|
|
|
|
from vllm.v1.worker.gpu_input_batch import InputBatch
|
|
|
|
DEFAULT_SAMPLING_PARAMS = dict(
|
|
temperature=-1.0,
|
|
min_p=0.0,
|
|
# strictly disabled for now
|
|
# top_k=-1,
|
|
# top_p=0.0,
|
|
# frequency_penalties=0.0,
|
|
# presence_penalties=0.0,
|
|
# repetition_penalties=0.0,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class TPUSupportedSamplingMetadata:
|
|
# This class exposes a more xla-friendly interface than SamplingMetadata
|
|
# on TPU, in particular all arguments should be traceable and no optionals
|
|
# are allowed, to avoid graph recompilation on Nones.
|
|
temperature: torch.Tensor
|
|
|
|
min_p: torch.Tensor
|
|
# Still too slow on forward_native!
|
|
top_k: torch.Tensor = None
|
|
top_p: torch.Tensor = None
|
|
|
|
# Greedy sampling flag for compiling single xla graph.
|
|
all_greedy: torch.Tensor = None
|
|
|
|
# Generator not supported by xla
|
|
generators: dict[int,
|
|
torch.Generator] = field(default_factory=lambda: dict())
|
|
|
|
# unsupported, you need to return an extra tensor of static size BxV
|
|
max_num_logprobs = None
|
|
|
|
# TODO No penalties for now
|
|
no_penalties: bool = True
|
|
prompt_token_ids = None
|
|
frequency_penalties = None
|
|
presence_penalties = None
|
|
repetition_penalties = None
|
|
# should use tensor
|
|
output_token_ids: list[list[int]] = field(default_factory=lambda: list())
|
|
|
|
min_tokens = None # impl is not vectorized
|
|
|
|
logit_bias: list[Optional[dict[int, float]]] = field(
|
|
default_factory=lambda: list())
|
|
|
|
allowed_token_ids_mask = None
|
|
bad_words_token_ids = None
|
|
indices_do_sample: torch.Tensor = None
|
|
|
|
@classmethod
|
|
def from_input_batch(
|
|
cls, input_batch: InputBatch,
|
|
indices_do_sample: torch.Tensor) -> "TPUSupportedSamplingMetadata":
|
|
"""
|
|
Copy sampling tensors slices from `input_batch` to on device tensors.
|
|
|
|
`InputBatch._make_sampling_metadata` causes recompilation on XLA as it
|
|
slices dynamic shapes on device tensors. This impl moves the dynamic
|
|
ops to CPU and produces tensors of fixed `padded_num_reqs` size. It
|
|
also reuses the on-device persistent tensors managed in `input_batch`
|
|
to reduce waste.
|
|
|
|
`indices_do_sample` contains the indices to be fed to the Sampler,
|
|
normally one per request, here padded to the closest pre-compiled shape
|
|
We expect sampling params tensors to be padded to the same fixed shape.
|
|
|
|
Eg. 3 requests, tensors padded to 4
|
|
temperature: [0.7, 0.2, 0.9]=>[0.7, 0.2, 0.9, 0.0]
|
|
sample indices: [4, 10, 11]=>indices_do_sample: [4, 10, 11, 0]
|
|
"""
|
|
num_reqs = input_batch.num_reqs
|
|
padded_num_reqs = len(indices_do_sample)
|
|
|
|
def copy_slice(cpu_tensor: torch.Tensor, tpu_tensor: torch.Tensor,
|
|
fill_val) -> torch.Tensor:
|
|
# Copy slice from CPU to corresponding TPU pre-allocated tensor.
|
|
# Pad value is the default one.
|
|
cpu_tensor[num_reqs:padded_num_reqs] = fill_val
|
|
tpu_tensor[:padded_num_reqs] = cpu_tensor[:padded_num_reqs]
|
|
|
|
# NOTE NickLucche The sync CPU-TPU graph we produce here must be
|
|
# consistent. We can't have flags to skip copies or we'll end up
|
|
# recompiling.
|
|
copy_slice(input_batch.temperature_cpu_tensor, input_batch.temperature,
|
|
DEFAULT_SAMPLING_PARAMS["temperature"])
|
|
# TODO Temporarily disabled until sampling options are enabled
|
|
# copy_slice(input_batch.top_p_cpu_tensor, input_batch.top_p)
|
|
# copy_slice(input_batch.top_k_cpu_tensor, input_batch.top_k)
|
|
copy_slice(input_batch.min_p_cpu_tensor, input_batch.min_p,
|
|
DEFAULT_SAMPLING_PARAMS["min_p"])
|
|
|
|
# copy_slice(input_batch.frequency_penalties_cpu_tensor,
|
|
# input_batch.frequency_penalties)
|
|
# copy_slice(input_batch.presence_penalties_cpu_tensor,
|
|
# input_batch.presence_penalties)
|
|
# copy_slice(input_batch.repetition_penalties_cpu_tensor,
|
|
# input_batch.repetition_penalties)
|
|
|
|
xm.mark_step()
|
|
xm.wait_device_ops()
|
|
|
|
# Slice persistent device tensors to a fixed pre-compiled padded shape.
|
|
return cls(
|
|
temperature=input_batch.temperature[:padded_num_reqs],
|
|
# Scalar tensor for xla-friendly tracing.
|
|
all_greedy=torch.tensor(input_batch.all_greedy,
|
|
dtype=torch.bool,
|
|
device=input_batch.device),
|
|
# TODO enable more and avoid returning None values
|
|
top_p=None, # input_batch.top_p[:padded_num_reqs],
|
|
top_k=None, # input_batch.top_k[:padded_num_reqs],
|
|
min_p=input_batch.min_p[:padded_num_reqs],
|
|
generators=input_batch.generators,
|
|
indices_do_sample=indices_do_sample)
|