lightseekorg--tokenspeed/python/tokenspeed/runtime/sampling/sampling_params.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Sampling parameters for text generation."""

import zlib
from typing import Any

_SAMPLING_EPS = 1e-6

# Sentinel for "top_k is disabled" (sample from whole vocab). We rewrite
# top_k=-1 (API convention) to this value so downstream code can pass it
# unchanged to top_k kernels that expect a positive cutoff.
_TOP_K_DISABLED = 1 << 30

# Upper bound the fused top-k + top-p kernel sorts in its on-chip top-K
# branch. Requests with a finite top_k above this would silently fall through
# to the top-p-only branch, so reject them at request time. Must stay in sync
# with K_TOPK_MAX in fused_topk_topp.h.
_TOP_K_FUSED_MAX = 128


class SamplingParams:
    """
    The sampling parameters.

    See docs/backend/sampling_params.md or
    https://docs.tokenspeed.ai/backend/sampling_params.html
    for the documentation.
    """

    def __init__(
        self,
        max_new_tokens: int | None = None,
        stop: str | list[str] | None = None,
        stop_token_ids: list[int] | None = None,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = -1,
        min_p: float = 0.0,
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        repetition_penalty: float = 1.0,
        min_new_tokens: int = 0,
        json_schema: str | None = None,
        regex: str | None = None,
        ebnf: str | None = None,
        structural_tag: str | None = None,
        ignore_eos: bool = False,
        skip_special_tokens: bool = True,
        spaces_between_special_tokens: bool = True,
        no_stop_trim: bool = False,
        thinking_budget: int | None = None,
        custom_params: dict[str, Any] | None = None,
        stream_interval: int | None = None,
        logit_bias: dict[str, float] | None = None,
        seed: int | None = None,
        # vLLM-style output logprobs. None = off; 0 = the sampled (generated)
        # token's logprob at each output position. Other values are rejected by
        # verify().
        logprobs: int | None = None,
        # OpenAI-compat: `n` is a request-level fanout (number of choices)
        # that the serving layer forwards on every sampling_params dict.
        # TokenSpeed does not multiplex a single request into n completions,
        # so accept and ignore.
        n: int = 1,
    ) -> None:
        self.max_new_tokens = max_new_tokens
        self.stop_strs = stop
        if stop_token_ids:
            self.stop_token_ids = set(stop_token_ids)
        else:
            self.stop_token_ids = None
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.min_p = min_p
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
        self.repetition_penalty = repetition_penalty
        self.min_new_tokens = min_new_tokens
        self.regex = regex
        self.json_schema = json_schema
        self.ebnf = ebnf
        self.structural_tag = structural_tag
        self.ignore_eos = ignore_eos
        self.skip_special_tokens = skip_special_tokens
        self.spaces_between_special_tokens = spaces_between_special_tokens
        self.no_stop_trim = no_stop_trim
        self.custom_params = custom_params
        self.thinking_budget = thinking_budget
        self.stream_interval = stream_interval
        self.logit_bias = logit_bias
        self.seed = seed
        self.logprobs = logprobs

        # Process some special cases
        if self.temperature < _SAMPLING_EPS:
            # top_k = 1 means greedy sampling
            self.temperature = 1.0
            self.top_k = 1
        if self.top_k == -1:
            self.top_k = _TOP_K_DISABLED

    def verify(self, vocab_size: int) -> None:
        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 not 0.0 <= self.min_p <= 1.0:
            raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.")
        if self.top_k < -1 or self.top_k == 0:
            raise ValueError(
                f"top_k must be -1 (disable), or at least 1, " f"got {self.top_k}."
            )
        if self.top_k != _TOP_K_DISABLED and self.top_k >= _TOP_K_FUSED_MAX:
            raise ValueError(
                f"top_k must be < {_TOP_K_FUSED_MAX} (fused kernel limit) "
                f"or -1 (disable), got {self.top_k}."
            )
        if not -2.0 <= self.frequency_penalty <= 2.0:
            raise ValueError(
                "frequency_penalty must be in [-2, 2], got "
                f"{self.frequency_penalty}."
            )
        if not -2.0 <= self.presence_penalty <= 2.0:
            raise ValueError(
                "presence_penalty must be in [-2, 2], got " f"{self.presence_penalty}."
            )
        if not 0.0 <= self.repetition_penalty <= 2.0:
            raise ValueError(
                "repetition_penalty must be in (0, 2], got "
                f"{self.repetition_penalty}."
            )
        if not 0 <= self.min_new_tokens:
            raise ValueError(
                f"min_new_tokens must be in (0, max_new_tokens], got "
                f"{self.min_new_tokens}."
            )
        if self.max_new_tokens is not None:
            if self.max_new_tokens < 0:
                raise ValueError(
                    f"max_new_tokens must be at least 0, got {self.max_new_tokens}."
                )
            if not self.min_new_tokens <= self.max_new_tokens:
                raise ValueError(
                    f"min_new_tokens must be in (0, max_new_tokens({self.max_new_tokens})], got "
                    f"{self.min_new_tokens}."
                )
        if self.logit_bias is not None:
            for token_id in self.logit_bias:
                if not 0 <= int(token_id) < vocab_size:
                    raise ValueError(
                        f"logit_bias must has keys in [0, {vocab_size - 1}], got "
                        f"{token_id}."
                    )

        if self.logprobs is not None and self.logprobs != 0:
            # Only the sampled token's logprob (logprobs=0) is materialized;
            # top-k (>0) and full-vocab (-1) output logprobs are not supported.
            raise ValueError(
                f"logprobs={self.logprobs} is not supported; use logprobs=0 "
                "(the sampled token's logprob)."
            )

        grammars = [
            self.json_schema,
            self.regex,
            self.ebnf,
        ]  # since mutually exclusive, only one can be set
        if sum(x is not None for x in grammars) > 1:
            raise ValueError("Only one of regex, json_schema, or ebnf can be set.")

    def requested_features(self) -> "set[str]":
        """Return the set of backend-facing feature names this request needs.

        `temperature`, `top_k`, `top_p`, `min_p` each appear only when the
        corresponding field is not at its neutral default. Used by
        SamplingBackend.register() to reject requests asking for features
        the active backend does not implement."""
        out: set[str] = set()
        if abs(self.temperature - 1.0) > _SAMPLING_EPS:
            out.add("temperature")
        # top_k=_TOP_K_DISABLED and top_k=1 (greedy short-circuit from __init__) are neutral.
        if self.top_k != _TOP_K_DISABLED and self.top_k != 1:
            out.add("top_k")
        if self.top_p < 1.0:
            out.add("top_p")
        if self.min_p > 0.0:
            out.add("min_p")
        if self.frequency_penalty != 0.0:
            out.add("frequency_penalty")
        if self.presence_penalty != 0.0:
            out.add("presence_penalty")
        if self.repetition_penalty != 1.0:
            out.add("repetition_penalty")
        if self.logit_bias:
            out.add("logit_bias")
        return out

    def resolve_seed(self, rid: str) -> None:
        """If the caller didn't supply a seed, derive one deterministically
        from rid. Called at the single request-materialization point so all
        TP/DP ranks agree on the seed."""
        if self.seed is None:
            self.seed = zlib.crc32(rid.encode("utf-8")) & 0xFFFFFFFF

    def normalize(self, tokenizer) -> None:
        # Process stop strings
        if self.stop_strs is None:
            self.stop_strs = []
            self.stop_str_max_len = 0
        else:
            if isinstance(self.stop_strs, str):
                self.stop_strs = [self.stop_strs]

            stop_str_max_len = 0
            for stop_str in self.stop_strs:
                if tokenizer is not None:
                    stop_str_ids = tokenizer.encode(stop_str, add_special_tokens=False)
                    stop_str_max_len = max(stop_str_max_len, len(stop_str_ids))
                else:
                    stop_str_max_len = max(stop_str_max_len, len(stop_str))
            self.stop_str_max_len = stop_str_max_len