# Copyright (c) 2023, The vLLM team.
# Copyright (c) 2023, FuriosaAI Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
from dataclasses import dataclass
from enum import IntEnum
from functools import cached_property
from pydantic import BaseModel
from furiosa_llm.constants import StructuredOutputsBackend as StructuredOutputsBackend
from furiosa_llm.metadata.tasks import PoolingTask
from furiosa_llm.outputs import RequestOutputKind
_SAMPLING_EPS = 1e-5
class SamplingType(IntEnum):
GREEDY = 0
RANDOM = 1
# https://github.com/vllm-project/vllm/blob/6d8d0a24c02bfd84d46b3016b865a44f048ae84b/vllm/sampling_params.py#L30-L84
@dataclass
class StructuredOutputsParams:
"""
Parameters for a pattern, corresponding to one state machine.
These parameters also function as request JSON body parameters (in "structured_outputs" field) in OpenAI-compatible interfaces.
"""
json: str | dict | None = None
regex: str | None = None
choice: list[str] | None = None
grammar: str | None = None
json_object: bool | None = None
disable_any_whitespace: bool = False
disable_additional_properties: bool = False
whitespace_pattern: str | None = None
structural_tag: str | None = None
# XXX: Indicates if this parameter was generated from xgrammar’s harmony built-in structural tag.
# This field is needed to work around the xgrammar limitation described in #3907.
# It should not be set by the user, but should be managed internally and implicitly.
_is_harmony_structural_tag: bool = False
backend: str | None = None
@staticmethod
def from_optional(
json: dict | BaseModel | str | None = None,
regex: str | None = None,
choice: list[str] | None = None,
grammar: str | None = None,
json_object: bool | None = None,
whitespace_pattern: str | None = None,
structural_tag: str | None = None,
backend: str | None = None,
disable_additional_properties: bool = False,
) -> "StructuredOutputsParams | None":
if all(arg is None for arg in (json, regex, choice, grammar, json_object, structural_tag)):
return None
# Extract json schemas from pydantic models
if isinstance(json, (BaseModel, type(BaseModel))):
json = json.model_json_schema()
return StructuredOutputsParams(
json=json,
regex=regex,
choice=choice,
grammar=grammar,
json_object=json_object,
whitespace_pattern=whitespace_pattern,
structural_tag=structural_tag,
disable_additional_properties=disable_additional_properties,
backend=backend,
)
def __post_init__(self):
"""Validate that some fields are mutually exclusive."""
guide_count = sum(
[
self.json is not None,
self.regex is not None,
self.choice is not None,
self.grammar is not None,
self.json_object is not None,
self.structural_tag is not None,
]
)
if guide_count > 1:
raise ValueError(
"You can only use one kind of structured output but multiple are "
f"specified: {self.__dict__}"
)
if guide_count < 1:
raise ValueError("at least one structured output constraint must be specified")
# https://github.com/vllm-project/vllm/blob/v0.6.2/vllm/sampling_params.py#L46-L127
[docs]
class SamplingParams:
"""Sampling parameters for text generation.
Args:
n: Number of output sequences to return for the given prompt.
repetition_penalty: Float that penalizes new tokens based on whether
they appear in the prompt and the generated text so far. Values > 1
encourage the model to use new tokens, while values < 1 encourage
the model to repeat tokens.
temperature: Float that controls the randomness of the sampling. Lower
values make the model more deterministic, while higher values make
the model more random. Zero means greedy sampling.
top_p: Float that controls the cumulative probability of the top tokens
to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
top_k: Integer that controls the number of top tokens to consider. Set
to -1 to consider all tokens.
min_p: Float that represents the minimum probability for a token to be
considered, relative to the probability of the most likely token.
Must be in [0, 1]. Set to 0 to disable this.
stop: List or single non-empty string(s) that stop generation when the
first string is matched. By default, the matched stop string is removed
from the returned text. When streaming, Furiosa-LLM may hold back
a trailing window up to the longest configured stop string to
prevent a partial match from being emitted.
stop_token_ids: Token IDs that stop the generation when they are generated. The returned
output will contain the stop tokens unless the stop tokens are special
tokens.
include_stop_str_in_output: Whether to retain the matched stop string
in the returned text. Defaults to False. When True, text after
the match within the token that completed it is removed. Streaming
does not need to hold back a trailing stop-string window in this
mode.
ignore_eos: Whether to ignore the EOS token and continue generating tokens
after the EOS token is generated.
max_tokens: Maximum number of tokens to generate per output sequence.
If the value is None, it is capped to the maximum sequence length.
min_tokens: Minimum number of tokens to generate per output sequence
before EOS or stop_token_ids can be generated
skip_special_tokens: Whether to skip special tokens in the output.
logprobs: Number of log probabilities to return per output token.
When set to None, no probability is returned. If set to a non-None
value, the result includes the log probabilities of the specified
number of most likely tokens, as well as the chosen tokens.
Note that the implementation follows the OpenAI API: The API will
always return the log probability of the sampled token, so there
may be up to `logprobs+1` elements in the response.
prompt_logprobs: Number of log probabilities to return per prompt token.
When set to None (default), no prompt logprobs are returned.
When set to a non-negative integer, returns the top-k log probabilities
for each prompt token position, plus the actual token's logprob.
Set to -1 to return log probabilities for all vocabulary tokens.
**Warning**: Using -1 can cause significant memory
and network overhead as it returns logprobs for the entire vocabulary
(e.g., ~150K tokens for Qwen models) at each prompt position.
structured_outputs: Parameter for specifying structured output requirements.
reasoning_structured_outputs: Parameter that matches the reasoning part of
model output. This will be used to strip the reasoning tokens before
applying structured output requirements on the answer content.
extra_args: Arbitrary additional arguments carried alongside the standard
sampling parameters. All extra (undeclared) keys of an API request are
collected here. Specific consumers extract the keys they care about;
e.g. ``kv_transfer_params`` is extracted from this dict.
kv_transfer_params: Access point of the remote KV connector, extracted
from ``extra_args["kv_transfer_params"]``. There can be one of two
ZMQ endpoints (``"host:port"``): ``"remote_decoder"`` and
``"remote_prefiller"``. If set, the worker connects to the given
endpoint to enable PD-disaggregated prefill or decoding. ``"id"`` is
the shared identifier of the P-D pair: it must be the same value on
both the prefiller and the decoder, and the connector uses it to
match the two sides.
"""
def __init__(
self,
*,
n: int = 1,
repetition_penalty: float = 1.0,
temperature: float = 1.0,
top_p: float = 1.0,
top_k: int = -1,
min_p: float = 0.0,
stop: str | list[str] | None = None,
stop_token_ids: list[int] | None = None,
include_stop_str_in_output: bool = False,
ignore_eos: bool = False,
max_tokens: int | None = 16,
min_tokens: int = 0,
skip_special_tokens: bool = True,
logprobs: int | None = None,
prompt_logprobs: int | None = None,
output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE,
structured_outputs: StructuredOutputsParams | None = None,
reasoning_structured_outputs: StructuredOutputsParams | None = None,
extra_args: dict | None = None,
) -> None:
# Filter empty strings and switch empty array to None
# This is to keep vLLM behavior compatibility
if isinstance(stop, str):
stop = [stop]
stop = [s for s in (stop or []) if s]
stop = stop or None
self.n = n
self.repetition_penalty = repetition_penalty
self.temperature = temperature
self.top_p = top_p
self.top_k = top_k
self.min_p = min_p
self.stop = stop
self.stop_token_ids = stop_token_ids
self.include_stop_str_in_output = include_stop_str_in_output
self.ignore_eos = ignore_eos
self.max_tokens = max_tokens
self.min_tokens = min_tokens
self.skip_special_tokens = skip_special_tokens
# https://github.com/vllm-project/vllm/blob/v0.6.2/vllm/sampling_params.py#L253
self.logprobs = 1 if logprobs is True else logprobs
self.prompt_logprobs = prompt_logprobs
self.output_kind = output_kind
self.structured_outputs = structured_outputs
self.reasoning_structured_outputs = reasoning_structured_outputs
self.extra_args = extra_args
self.kv_transfer_params = (extra_args or {}).get("kv_transfer_params")
self._verify_args()
@classmethod
def from_optional(
cls,
*,
n: int | None = None,
repetition_penalty: float | None = 1.0,
temperature: float | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float = 0.0,
stop: str | list[str] | None = None,
stop_token_ids: list[int] | None = None,
include_stop_str_in_output: bool = False,
ignore_eos: bool | None = None,
max_tokens: int | None = None,
min_tokens: int | None = None,
skip_special_tokens: bool = True,
logprobs: int | None = None,
prompt_logprobs: int | None = None,
output_kind: RequestOutputKind | None = None,
structured_outputs: StructuredOutputsParams | None = None,
reasoning_structured_outputs: StructuredOutputsParams | None = None,
extra_args: dict | None = None,
) -> "SamplingParams":
return cls(
n=1 if n is None else n,
repetition_penalty=1.0 if repetition_penalty is None else repetition_penalty,
temperature=1.0 if temperature is None else temperature,
top_p=1.0 if top_p is None else top_p,
top_k=-1 if top_k is None else top_k,
min_p=min_p,
stop=stop,
stop_token_ids=stop_token_ids,
include_stop_str_in_output=include_stop_str_in_output,
ignore_eos=False if ignore_eos is None else ignore_eos,
max_tokens=max_tokens,
min_tokens=0 if min_tokens is None else min_tokens,
skip_special_tokens=skip_special_tokens,
logprobs=logprobs,
prompt_logprobs=prompt_logprobs,
output_kind=RequestOutputKind.CUMULATIVE if output_kind is None else output_kind,
structured_outputs=structured_outputs,
reasoning_structured_outputs=reasoning_structured_outputs,
extra_args=extra_args,
)
def clone(self) -> "SamplingParams":
return copy.deepcopy(self)
def __eq__(self, other) -> bool:
return (
isinstance(other, SamplingParams)
and self.n == other.n
and self.repetition_penalty == other.repetition_penalty
and self.temperature == other.temperature
and self.top_p == other.top_p
and self.top_k == other.top_k
and self.min_p == other.min_p
and self.max_tokens == other.max_tokens
and self.min_tokens == other.min_tokens
and self.skip_special_tokens == other.skip_special_tokens
and self.logprobs == other.logprobs
and self.prompt_logprobs == other.prompt_logprobs
and self.output_kind == other.output_kind
and self.stop == other.stop
and self.stop_token_ids == other.stop_token_ids
and self.include_stop_str_in_output == other.include_stop_str_in_output
and self.ignore_eos == other.ignore_eos
and self.structured_outputs == other.structured_outputs
and self.reasoning_structured_outputs == other.reasoning_structured_outputs
)
def _verify_args(self) -> None:
if self.n < 1:
raise ValueError(f"n must be at least 1, got {self.n}.")
if self.n > 1:
raise ValueError(f"furiosa-llm currently does not support n > 1, got {self.n}.")
if not 0.0 < self.repetition_penalty <= 2.0:
raise ValueError(
f"repetition_penalty must be in (0, 2], got {self.repetition_penalty}."
)
if self.temperature < 0.0:
raise ValueError(f"temperature must be non-negative, got {self.temperature}.")
if not 0.0 < self.top_p <= 1.0:
raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
if self.top_k < -1 or self.top_k == 0:
raise ValueError(f"top_k must be -1 (disable), or at least 1, got {self.top_k}.")
if not 0.0 <= self.min_p <= 1.0:
raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.")
if self.stop is not None:
if any(len(s) == 0 for s in self.stop):
raise ValueError("stop strings must not be empty.")
if self.max_tokens is not None and self.max_tokens < 1:
raise ValueError(f"max_tokens must be at least 1, got {self.max_tokens}.")
if self.min_tokens < 0:
raise ValueError(
f"min_tokens must be greater than or equal to 0, got {self.min_tokens}."
)
if self.max_tokens is not None and self.min_tokens > self.max_tokens:
raise ValueError(
f"min_tokens must be less than or equal to "
f"max_tokens={self.max_tokens}, got {self.min_tokens}."
)
if self.logprobs is not None and self.logprobs < 0:
raise ValueError(f"logprobs must be non-negative, got {self.logprobs}.")
if (
self.prompt_logprobs is not None
and self.prompt_logprobs != -1
and self.prompt_logprobs < 0
):
raise ValueError(
f"prompt_logprobs must be -1 (all tokens) or non-negative, "
f"got {self.prompt_logprobs}."
)
if self.kv_transfer_params is not None:
if not isinstance(self.kv_transfer_params, dict):
raise ValueError(
f"kv_transfer_params must be a dict or None, got {self.kv_transfer_params}."
)
if "id" not in self.kv_transfer_params:
raise ValueError(
f"kv_transfer_params must contain an 'id' field, got {self.kv_transfer_params}."
)
if (
"remote_prefiller" not in self.kv_transfer_params
and "remote_decoder" not in self.kv_transfer_params
):
raise ValueError(
f"kv_transfer_params must contain either 'remote_prefiller' or 'remote_decoder' field, got {self.kv_transfer_params}."
)
def structured_outputs_enabled(self):
return self.structured_outputs is not None
@cached_property
def sampling_type(self) -> SamplingType:
if self.temperature < _SAMPLING_EPS:
return SamplingType.GREEDY
return SamplingType.RANDOM
def __repr__(self) -> str:
return (
f"SamplingParams(n={self.n}, "
f"repetition_penalty={self.repetition_penalty}, "
f"temperature={self.temperature}, "
f"top_p={self.top_p}, "
f"top_k={self.top_k}, "
f"min_p={self.min_p}, "
f"max_tokens={self.max_tokens}, "
f"min_tokens={self.min_tokens}, "
f"skip_special_tokens={self.skip_special_tokens}, "
f"logprobs={self.logprobs}, "
f"prompt_logprobs={self.prompt_logprobs}, "
f"structured_outputs={self.structured_outputs}, "
f"reasoning_structured_outputs={self.reasoning_structured_outputs}, "
f"output_kind={self.output_kind}, "
f"stop={self.stop}, "
f"stop_token_ids={self.stop_token_ids}, "
f"include_stop_str_in_output={self.include_stop_str_in_output}, "
f"ignore_eos={self.ignore_eos}, "
f"extra_args={self.extra_args})"
)
def validate_speculative_sampling_params(
sampling_params: SamplingParams, speculative_decoding_enabled: bool
) -> None:
"""Reject request features not yet supported by speculative decoding."""
if not speculative_decoding_enabled:
return
if sampling_params.kv_transfer_params is not None:
raise ValueError("PD disaggregation is not supported with speculative decoding")
if (
sampling_params.structured_outputs is not None
or sampling_params.reasoning_structured_outputs is not None
):
raise ValueError("structured outputs are not supported with speculative decoding")
if sampling_params.logprobs is not None or sampling_params.prompt_logprobs is not None:
raise ValueError("logprobs and prompt_logprobs are not supported with speculative decoding")
[docs]
class PoolingParams:
"""
API parameters for pooling models.
Attributes:
truncate_prompt_tokens: Controls prompt truncation.
Set to -1 to use the model's default truncation size.
Set to k to keep only the last k tokens (left truncation).
Set to None to disable truncation.
dimensions: Number of dimensions for the output embedding.
If set, truncates the embedding to the first N dimensions
(Matryoshka Representation Learning). Must be a positive integer.
normalize: Whether to normalize the embeddings outputs. Only supported for embedding tasks.
"""
def __init__(
self,
truncate_prompt_tokens: int | None = None,
dimensions: int | None = None,
normalize: bool | None = True,
task: PoolingTask | None = None,
):
if truncate_prompt_tokens is not None and truncate_prompt_tokens < -1:
raise ValueError(
f"truncate_prompt_tokens must be -1, None, or a non-negative integer, "
f"got {truncate_prompt_tokens}."
)
# TODO: check if the model supports MRL before allowing dimension truncation
if dimensions is not None and dimensions < 1:
raise ValueError(f"dimensions must be a positive integer, got {dimensions}")
self.truncate_prompt_tokens = truncate_prompt_tokens
self.dimensions = dimensions
self.normalize = normalize
self.task = task
