llm.generation — Sampling and Backends

Token sampling and the generation backend abstraction. The backend is what the serving tier and the trainer's evaluation loop both call into.

Sampling Utilities

sampling

Shared token sampling utilities for generation backends.

apply_repetition_penalty

apply_repetition_penalty(logits, token_ids, repetition_penalty)

Apply repetition penalty in-place on 1D logits.

Token ids outside [0, vocab_size) are silently dropped, matching the other penalty helpers (:func:apply_frequency_penalty, :func:apply_presence_penalty, :func:apply_logit_bias). Without this the torch.gather below raises an out-of-bounds error for any id that is not representable in these logits (e.g. a truncation or API boundary passing ids the model's vocabulary never produced).

源代码位于： src/llm/generation/sampling.py

def apply_repetition_penalty(
    logits: torch.Tensor,
    token_ids: list[int],
    repetition_penalty: float,
) -> torch.Tensor:
    """Apply repetition penalty in-place on 1D logits.

    Token ids outside ``[0, vocab_size)`` are silently dropped, matching
    the other penalty helpers (:func:`apply_frequency_penalty`,
    :func:`apply_presence_penalty`, :func:`apply_logit_bias`). Without
    this the ``torch.gather`` below raises an out-of-bounds error for any
    id that is not representable in these logits (e.g. a truncation or
    API boundary passing ids the model's vocabulary never produced).
    """
    if repetition_penalty == 1.0 or not token_ids:
        return logits
    # ``repetition_penalty <= 0`` silently corrupts the distribution: 0 divides
    # by zero (inf) and a negative value flips each seen score's sign, producing
    # non-finite or inverted logits. Values in (0, 1) legitimately *encourage*
    # repetition, so only reject <= 0 (RIL TASK-250).
    if repetition_penalty <= 0:
        raise ValueError(f"repetition_penalty must be > 0, got {repetition_penalty!r}")

    vocab_size = logits.size(-1)
    valid_ids = [tid for tid in token_ids if 0 <= tid < vocab_size]
    if not valid_ids:
        return logits

    adjusted = logits.clone()
    device = adjusted.device
    ids = torch.tensor(valid_ids, device=device)
    scores = torch.gather(adjusted, 0, ids)
    scores = torch.where(scores < 0, scores * repetition_penalty, scores / repetition_penalty)
    adjusted.scatter_(0, ids, scores)
    return adjusted

apply_frequency_penalty

apply_frequency_penalty(logits, token_ids, frequency_penalty)

Subtract frequency_penalty * count(token) from each seen token's logit.

Implements the OpenAI-compatible frequency_penalty semantics (see https://platform.openai.com/docs/api-reference/chat/create): positive values penalise tokens in proportion to how often they have already appeared in the generated text. Zero (the default) is a no-op so callers don't need to special-case the off state.

参数：

名称	类型	描述	默认
logits	Tensor	1D [vocab_size] tensor. Not mutated.	必需
token_ids	list[int]	List of token ids generated so far (may include duplicates; duplicates count toward the penalty).	必需
frequency_penalty	float	Penalty coefficient. 0.0 disables the adjustment; values typically live in [-2.0, 2.0].	必需

返回：

类型	描述
Tensor	A new 1D tensor with the per-frequency penalty subtracted.

源代码位于： src/llm/generation/sampling.py

def apply_frequency_penalty(
    logits: torch.Tensor,
    token_ids: list[int],
    frequency_penalty: float,
) -> torch.Tensor:
    """Subtract ``frequency_penalty * count(token)`` from each seen token's logit.

    Implements the OpenAI-compatible ``frequency_penalty`` semantics
    (see https://platform.openai.com/docs/api-reference/chat/create):
    positive values penalise tokens in proportion to how often they
    have already appeared in the generated text. Zero (the default)
    is a no-op so callers don't need to special-case the off state.

    Args:
        logits: 1D ``[vocab_size]`` tensor. Not mutated.
        token_ids: List of token ids generated so far (may include
            duplicates; duplicates count toward the penalty).
        frequency_penalty: Penalty coefficient. ``0.0`` disables the
            adjustment; values typically live in ``[-2.0, 2.0]``.

    Returns:
        A new 1D tensor with the per-frequency penalty subtracted.
    """
    if frequency_penalty == 0.0 or not token_ids:
        return logits

    counts = Counter(token_ids)
    vocab_size = logits.size(-1)
    # Drop ids that fall outside the vocab — they're not representable
    # in these logits, so penalising them is meaningless and would
    # raise an index error on the scatter below.
    valid_ids = {tid: c for tid, c in counts.items() if 0 <= tid < vocab_size}
    if not valid_ids:
        return logits

    adjusted = logits.clone()
    device = adjusted.device
    ids = torch.tensor(list(valid_ids), device=device, dtype=torch.long)
    penalties = torch.tensor([valid_ids[tid] for tid in valid_ids], device=device, dtype=adjusted.dtype)
    adjusted.scatter_add_(
        0,
        ids,
        -frequency_penalty * penalties,
    )
    return adjusted

apply_presence_penalty

apply_presence_penalty(logits, token_ids, presence_penalty)

Subtract a flat presence_penalty from each seen token's logit.

Implements the OpenAI-compatible presence_penalty semantics (see https://platform.openai.com/docs/api-reference/chat/create): positive values penalise tokens that have appeared at least once in the generated text, encouraging the model to talk about new topics. The penalty is flat — a token that appeared 5 times is penalised the same as one that appeared once. That is the key distinction from :func:apply_frequency_penalty, which scales by count.

Negative values boost seen tokens (less common, but valid per OpenAI's spec — useful when you want the model to stay on topic).

参数：

名称	类型	描述	默认
logits	Tensor	1D [vocab_size] tensor. Not mutated.	必需
token_ids	list[int]	List of token ids generated so far. Order and duplicates are ignored — only the set matters.	必需
presence_penalty	float	Penalty coefficient. 0.0 is a no-op; values typically live in [-2.0, 2.0].	必需

返回：

类型	描述
Tensor	A new 1D tensor with the flat per-presence penalty applied.

源代码位于： src/llm/generation/sampling.py

def apply_presence_penalty(
    logits: torch.Tensor,
    token_ids: list[int],
    presence_penalty: float,
) -> torch.Tensor:
    """Subtract a flat ``presence_penalty`` from each **seen** token's logit.

    Implements the OpenAI-compatible ``presence_penalty`` semantics
    (see https://platform.openai.com/docs/api-reference/chat/create):
    positive values penalise tokens that have appeared **at least
    once** in the generated text, encouraging the model to talk
    about new topics. The penalty is **flat** — a token that
    appeared 5 times is penalised the same as one that appeared
    once. That is the key distinction from
    :func:`apply_frequency_penalty`, which scales by count.

    Negative values *boost* seen tokens (less common, but valid per
    OpenAI's spec — useful when you want the model to stay on
    topic).

    Args:
        logits: 1D ``[vocab_size]`` tensor. Not mutated.
        token_ids: List of token ids generated so far. Order and
            duplicates are ignored — only the **set** matters.
        presence_penalty: Penalty coefficient. ``0.0`` is a no-op;
            values typically live in ``[-2.0, 2.0]``.

    Returns:
        A new 1D tensor with the flat per-presence penalty applied.
    """
    if presence_penalty == 0.0 or not token_ids:
        return logits

    vocab_size = logits.size(-1)
    # Only the set of seen ids matters, not the counts.
    seen = {tid for tid in token_ids if 0 <= tid < vocab_size}
    if not seen:
        return logits

    adjusted = logits.clone()
    device = adjusted.device
    ids = torch.tensor(list(seen), device=device, dtype=torch.long)
    adjusted.scatter_add_(
        0,
        ids,
        -presence_penalty * torch.ones(len(seen), device=device, dtype=adjusted.dtype),
    )
    return adjusted

apply_logit_bias

apply_logit_bias(logits, logit_bias)

Add a per-token additive bias to 1D logits before sampling.

Implements the OpenAI-compatible logit_bias semantics (see https://platform.openai.com/docs/api-reference/chat/create): the bias is added to the affected token's logit prior to sampling. Negative values discourage the token (down to -100 for a hard ban in OpenAI's spec); positive values encourage it (up to +100 for near-exclusive selection).

The bias is applied after the penalty helpers (:func:apply_repetition_penalty, :func:apply_frequency_penalty, :func:apply_presence_penalty). Rationale: a penalty subtracts to discourage repetition, and the bias is a user-intent override — applying it last lets the bias dominate any natural penalty the model would otherwise impose. This matches OpenAI's reference ordering (logit-bias is the final logit-stage modification before sampling).

参数：

名称	类型	描述	默认
logits	Tensor	1D [vocab_size] tensor. Not mutated.	必需
logit_bias	Mapping[Any, float] | None	Mapping {token_id: bias} to add. Keys may be int (internal use) or str (JSON boundary — OpenAI's spec uses string keys because JSON object keys are always strings). String keys are coerced via int() and invalid entries are silently dropped. None or empty disables the adjustment.	必需

返回：

类型	描述
Tensor	A new 1D tensor with the per-token bias added.

源代码位于： src/llm/generation/sampling.py

def apply_logit_bias(
    logits: torch.Tensor,
    logit_bias: Mapping[Any, float] | None,
) -> torch.Tensor:
    """Add a per-token additive bias to 1D logits before sampling.

    Implements the OpenAI-compatible ``logit_bias`` semantics
    (see https://platform.openai.com/docs/api-reference/chat/create):
    the bias is added to the affected token's logit prior to
    sampling. Negative values discourage the token (down to ``-100``
    for a hard ban in OpenAI's spec); positive values encourage it
    (up to ``+100`` for near-exclusive selection).

    The bias is applied **after** the penalty helpers
    (:func:`apply_repetition_penalty`,
    :func:`apply_frequency_penalty`,
    :func:`apply_presence_penalty`). Rationale: a penalty subtracts
    to discourage repetition, and the bias is a user-intent override
    — applying it last lets the bias dominate any natural penalty
    the model would otherwise impose. This matches OpenAI's
    reference ordering (logit-bias is the final logit-stage
    modification before sampling).

    Args:
        logits: 1D ``[vocab_size]`` tensor. Not mutated.
        logit_bias: Mapping ``{token_id: bias}`` to add. Keys may
            be ``int`` (internal use) or ``str`` (JSON boundary —
            OpenAI's spec uses string keys because JSON object keys
            are always strings). String keys are coerced via
            ``int()`` and invalid entries are silently dropped.
            ``None`` or empty disables the adjustment.

    Returns:
        A new 1D tensor with the per-token bias added.
    """
    if not logit_bias:
        return logits

    vocab_size = logits.size(-1)
    # Drop ids that fall outside the vocab — they're not
    # representable in these logits, so biasing them is meaningless
    # and would raise an index error on the scatter below. Coerce
    # str→int for the JSON-boundary case (OpenAI's spec uses string
    # keys because JSON object keys are always strings).
    valid: dict[int, float] = {}
    for tid, bias in logit_bias.items():
        try:
            tid_int = int(tid)
        except TypeError, ValueError:
            continue
        if 0 <= tid_int < vocab_size:
            valid[tid_int] = float(bias)
    if not valid:
        return logits

    adjusted = logits.clone()
    device = adjusted.device
    ids = torch.tensor(list(valid), device=device, dtype=torch.long)
    biases = torch.tensor(list(valid.values()), device=device, dtype=adjusted.dtype)
    adjusted.index_add_(0, ids, biases)
    return adjusted

mask_undecodable_logits

mask_undecodable_logits(logits, tokenizer_vocab_size)

Mask every logit whose token id the tokenizer cannot decode.

sample_next_token returns any id in [0, model_vocab); when the model's vocabulary is larger than the tokenizer's (a padded vocab, or a BPE/HF model served with a char tokenizer), sampling a tail id used to crash mid-stream — tokenizer.decode([token_id]) raises KeyError after part of the text was already yielded (RIL ISS-125). The penalty helpers already guard out-of-range ids; the sampled id itself must be bounded to the tokenizer's decodeable range too, by zeroing the tail probability mass (equivalently pinning those logits to -inf).

参数：

名称	类型	描述	默认
logits	Tensor	1D [vocab_size] or [B, vocab_size] logits. Mutated in place.	必需
tokenizer_vocab_size	int | None	The tokenizer's decodeable vocabulary size. None or >= the logits' vocab means nothing to mask.	必需

源代码位于： src/llm/generation/sampling.py

def mask_undecodable_logits(logits: torch.Tensor, tokenizer_vocab_size: int | None) -> None:
    """Mask every logit whose token id the tokenizer cannot decode.

    ``sample_next_token`` returns any id in ``[0, model_vocab)``; when the
    model's vocabulary is *larger* than the tokenizer's (a padded vocab, or a
    BPE/HF model served with a char tokenizer), sampling a tail id used to
    crash mid-stream — ``tokenizer.decode([token_id])`` raises ``KeyError``
    after part of the text was already yielded (RIL ISS-125). The penalty
    helpers already guard out-of-range ids; the sampled id itself must be
    bounded to the tokenizer's decodeable range too, by zeroing the tail
    probability mass (equivalently pinning those logits to ``-inf``).

    Args:
        logits: 1D ``[vocab_size]`` or ``[B, vocab_size]`` logits. Mutated
            in place.
        tokenizer_vocab_size: The tokenizer's decodeable vocabulary size.
            ``None`` or ``>=`` the logits' vocab means nothing to mask.
    """
    if tokenizer_vocab_size is None:
        return
    vocab_size = logits.size(-1)
    if tokenizer_vocab_size >= vocab_size:
        return
    if logits.dim() == 1:
        logits[tokenizer_vocab_size:] = -float("inf")
    else:
        logits[:, tokenizer_vocab_size:] = -float("inf")

sampling_probs

sampling_probs(logits, *, temperature=1.0, top_k=None, top_p=None)

Return the exact softmax distribution sample_next_token draws from.

Applies temperature scaling, top-k and top-p filtering identically to :func:sample_next_token and returns the full-vocab probability vector (masked-out tokens carry zero mass). Used by speculative decoding to score acceptance ratios against the same filtered distributions the draft/target samplers actually propose from — scoring against the raw full-vocab softmax would make the accepted-token set diverge from the eager backend's output (RIL ISS-99).

参数：

名称	类型	描述	默认
logits	Tensor	1D [vocab_size] logits. Not mutated.	必需
temperature	float	Sampling temperature. Must be non-zero — the caller handles the temperature == 0 (greedy) case via argmax.	1.0
top_k	int | None	Top-k filter; only the top_k largest logits survive.	None
top_p	float | None	Nucleus filter; the smallest tokens whose cumulative probability exceeds top_p are masked out.	None

源代码位于： src/llm/generation/sampling.py

def sampling_probs(
    logits: torch.Tensor,
    *,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
) -> torch.Tensor:
    """Return the exact softmax distribution ``sample_next_token`` draws from.

    Applies temperature scaling, top-k and top-p filtering identically to
    :func:`sample_next_token` and returns the full-vocab probability vector
    (masked-out tokens carry zero mass). Used by speculative decoding to
    score acceptance ratios against the **same filtered distributions** the
    draft/target samplers actually propose from — scoring against the raw
    full-vocab softmax would make the accepted-token set diverge from the
    eager backend's output (RIL ISS-99).

    Args:
        logits: 1D ``[vocab_size]`` logits. Not mutated.
        temperature: Sampling temperature. Must be non-zero — the
            caller handles the ``temperature == 0`` (greedy) case via
            ``argmax``.
        top_k: Top-k filter; only the ``top_k`` largest logits survive.
        top_p: Nucleus filter; the smallest tokens whose cumulative
            probability exceeds ``top_p`` are masked out.
    """
    # A non-positive temperature would silently invert the logits (anti-greedy)
    # or divide by zero. ``temperature == 0`` is the caller's greedy case
    # (handled in :func:`sample_next_token` before this is reached), so any
    # value here must be positive — reject instead of emitting a corrupt
    # distribution (RIL TASK-249).
    if temperature <= 0:
        raise ValueError(f"temperature must be > 0, got {temperature!r}")
    # ``top_k <= 0`` makes torch.topk's ``k`` out of range and surfaces as a new
    # cryptic IndexError; reject it with a clear message (RIL TASK-250).
    if top_k is not None and top_k < 1:
        raise ValueError(f"top_k must be >= 1 when set, got {top_k!r}")

    next_logits = logits / temperature

    if top_k is not None:
        vocab_size = next_logits.size(-1)
        values, _ = torch.topk(next_logits, min(top_k, vocab_size))
        next_logits = next_logits.clone()
        next_logits[next_logits < values[-1]] = -torch.inf

    if top_p is not None and 0.0 < top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
        sorted_indices_to_remove[0] = False
        next_logits = next_logits.clone()
        next_logits[sorted_indices[sorted_indices_to_remove]] = -float("inf")

    return torch.softmax(next_logits, dim=-1)

sample_next_token

sample_next_token(logits, *, temperature=1.0, top_k=None, top_p=None)

Sample one token id from 1D logits.

源代码位于： src/llm/generation/sampling.py

def sample_next_token(
    logits: torch.Tensor,
    *,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
) -> int:
    """Sample one token id from 1D logits."""
    if temperature == 0:
        return int(torch.argmax(logits, dim=-1).item())
    if temperature < 0:
        # A negative temperature silently inverts the logits, making the
        # sampler anti-greedy (draw the lowest-logit token) with no error.
        # Reject it so silent generation corruption is impossible (RIL TASK-249).
        raise ValueError(f"temperature must be >= 0, got {temperature!r}")

    probs = sampling_probs(
        logits,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
    )
    return int(torch.multinomial(probs, num_samples=1).item())

Generation Backend ABC

backends

Generation backend abstractions.

GenerationConfig dataclass

Shared generation hyperparameters across inference backends.

源代码位于： src/llm/generation/backends.py

@dataclass(frozen=True)
class GenerationConfig:
    """Shared generation hyperparameters across inference backends."""

    max_new_tokens: int = 128
    temperature: float = 1.0
    top_k: int | None = None
    top_p: float | None = None
    repetition_penalty: float = 1.0
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0
    logit_bias: dict[int, float] | None = None
    use_cache: bool = True
    # OpenAI-compat ``stop``: generation halts the moment the streamed
    # output contains any of these as a suffix; the stop string itself
    # is NOT included in the final response. Accepts a single string or
    # a list of up to 4 strings (OpenAI's documented cap). None means
    # no stop — generation runs to ``max_new_tokens`` (default).
    stop: str | list[str] | None = None

GenerationBackend

Bases: ABC

Backend protocol for text generation.

源代码位于： src/llm/generation/backends.py

class GenerationBackend(abc.ABC):
    """Backend protocol for text generation."""

    @abc.abstractmethod
    def stream(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
        *,
        request_id: str | None = None,
    ) -> Generator[str]:
        pass

    def generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
        *,
        request_id: str | None = None,
    ) -> str:
        chunks = list(self.stream(model, tokenizer, prompt, config, request_id=request_id))
        return prompt + "".join(chunks)

    def batch_generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompts: list[str],
        config: GenerationConfig,
    ) -> list[str]:
        return [self.generate(model, tokenizer, prompt, config) for prompt in prompts]

EagerGenerationBackend

Bases: GenerationBackend

Default in-process generation using the library stream_generate path.

源代码位于： src/llm/generation/backends.py

class EagerGenerationBackend(GenerationBackend):
    """Default in-process generation using the library stream_generate path."""

    def stream(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
        *,
        request_id: str | None = None,  # accepted for the ABC; eager has no request identity
    ) -> Generator[str]:
        from llm.generation.eager import stream_generate

        yield from stream_generate(
            model=model,
            tokenizer=tokenizer,
            prompt=prompt,
            max_new_tokens=config.max_new_tokens,
            temperature=config.temperature,
            top_k=config.top_k,
            top_p=config.top_p,
            repetition_penalty=config.repetition_penalty,
            frequency_penalty=config.frequency_penalty,
            presence_penalty=config.presence_penalty,
            logit_bias=config.logit_bias,
            use_cache=config.use_cache,
            stop=config.stop,
        )

    def batch_generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompts: list[str],
        config: GenerationConfig,
    ) -> list[str]:
        from llm.generation.eager import batch_generate

        return batch_generate(
            model=model,
            tokenizer=tokenizer,
            prompts=prompts,
            max_new_tokens=config.max_new_tokens,
            temperature=config.temperature,
            top_k=config.top_k,
            top_p=config.top_p,
            repetition_penalty=config.repetition_penalty,
            frequency_penalty=config.frequency_penalty,
            presence_penalty=config.presence_penalty,
            logit_bias=cast(Any, config.logit_bias),
            stop=config.stop,
        )

BatchedGenerationBackend

Bases: GenerationBackend

Generation via ContinuousBatchingEngine (iteration-level scheduling).

源代码位于： src/llm/generation/backends.py

class BatchedGenerationBackend(GenerationBackend):
    """Generation via ContinuousBatchingEngine (iteration-level scheduling)."""

    def __init__(self, engine: ContinuousBatchingEngine):
        self.engine = engine

    def stream(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
        *,
        request_id: str | None = None,
    ) -> Generator[str]:
        from llm.serving.schemas import GenerationRequest

        request = GenerationRequest(
            request_id=request_id,
            prompt=prompt,
            max_new_tokens=config.max_new_tokens,
            temperature=config.temperature,
            top_k=config.top_k,
            top_p=config.top_p,
            repetition_penalty=config.repetition_penalty,
            frequency_penalty=config.frequency_penalty,
            presence_penalty=config.presence_penalty,
            logit_bias=cast(Any, config.logit_bias),
            stop=config.stop,
        )
        yield from self.engine.stream_request(request)

    def batch_generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompts: list[str],
        config: GenerationConfig,
    ) -> list[str]:
        from llm.serving.schemas import GenerationRequest

        requests = [
            GenerationRequest(
                prompt=prompt,
                max_new_tokens=config.max_new_tokens,
                temperature=config.temperature,
                top_k=config.top_k,
                top_p=config.top_p,
                repetition_penalty=config.repetition_penalty,
                frequency_penalty=config.frequency_penalty,
                presence_penalty=config.presence_penalty,
                logit_bias=cast(Any, config.logit_bias),
                stop=config.stop,
            )
            for prompt in prompts
        ]
        return self.engine.batch_generate_requests(requests)

SpeculativeDecodingBackend

Bases: GenerationBackend

Speculative decoding: small draft model proposes, large target verifies.

Implements Leviathan et al. 2023 - the draft model speculates gamma tokens ahead; the target scores them in a single forward pass and accepts each with probability min(1, q_target / q_draft). On rejection, sample a correction token from (q_target - q_draft)+. The output distribution exactly matches the target distribution under the same sampling parameters.

The model argument to :meth:stream / :meth:batch_generate is ignored - the target and draft models are bound at construction time. tokenizer must be the shared tokenizer used by both models (same vocab, pad id, eos id).

参数：

名称	类型	描述	默认
target_model	DecoderModel	The "expensive" model whose distribution is the canonical output distribution.	必需
draft_model	DecoderModel	The "cheap" model used for speculation. Must share vocabulary with target_model.	必需
gamma	int	Number of speculative tokens per round (default 5). Typical values: 4-8.	5

源代码位于： src/llm/generation/backends.py

class SpeculativeDecodingBackend(GenerationBackend):
    """Speculative decoding: small draft model proposes, large target verifies.

    Implements Leviathan et al. 2023 - the draft model speculates
    ``gamma`` tokens ahead; the target scores them in a single
    forward pass and accepts each with probability
    ``min(1, q_target / q_draft)``. On rejection, sample a
    correction token from ``(q_target - q_draft)+``. The output
    distribution exactly matches the target distribution under the
    same sampling parameters.

    The ``model`` argument to :meth:`stream` / :meth:`batch_generate`
    is **ignored** - the target and draft models are bound at
    construction time. ``tokenizer`` must be the shared tokenizer
    used by both models (same vocab, pad id, eos id).

    Args:
        target_model: The "expensive" model whose distribution is
            the canonical output distribution.
        draft_model: The "cheap" model used for speculation. Must
            share vocabulary with ``target_model``.
        gamma: Number of speculative tokens per round (default 5).
            Typical values: 4-8.
    """

    def __init__(
        self,
        target_model: DecoderModel,
        draft_model: DecoderModel,
        *,
        gamma: int = 5,
    ) -> None:
        if gamma < 1:
            raise ValueError(f"gamma must be >= 1, got {gamma}")
        self.target_model = target_model
        self.draft_model = draft_model
        self.gamma = gamma

    def stream(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
        *,
        request_id: str | None = None,  # accepted for the ABC; speculative has no request identity
    ) -> Generator[str]:
        from llm.generation.speculative import speculative_generate

        yield from speculative_generate(
            target=self.target_model,
            draft=self.draft_model,
            tokenizer=tokenizer,
            prompt=prompt,
            max_new_tokens=config.max_new_tokens,
            gamma=self.gamma,
            temperature=config.temperature,
            top_k=config.top_k,
            top_p=config.top_p,
            repetition_penalty=config.repetition_penalty,
            frequency_penalty=config.frequency_penalty,
            presence_penalty=config.presence_penalty,
            logit_bias=cast(Any, config.logit_bias),
            stop=config.stop,
        )

    def batch_generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompts: list[str],
        config: GenerationConfig,
    ) -> list[str]:
        return [self.generate(model, tokenizer, prompt, config) for prompt in prompts]

Backend Registry

registry

Generation backend registry and bootstrap.

build_speculative_backend

build_speculative_backend(*, target_model=None, draft_model=None, gamma=5, **_kwargs)

Build a speculative decoding backend (Leviathan et al., 2023).

Both target_model and draft_model must share vocabulary with the tokenizer passed at generation time. The gamma parameter controls how many candidate tokens the draft proposes per round.

源代码位于： src/llm/generation/registry.py

def build_speculative_backend(
    *,
    target_model: DecoderModel | None = None,
    draft_model: DecoderModel | None = None,
    gamma: int = 5,
    **_kwargs: Any,
) -> GenerationBackend:
    """Build a speculative decoding backend (Leviathan et al., 2023).

    Both ``target_model`` and ``draft_model`` must share vocabulary
    with the tokenizer passed at generation time. The ``gamma``
    parameter controls how many candidate tokens the draft proposes
    per round.
    """
    from llm.generation.backends import SpeculativeDecodingBackend

    if target_model is None or draft_model is None:
        raise ValueError("speculative backend requires both target_model and draft_model kwargs")
    return SpeculativeDecodingBackend(
        target_model=target_model,
        draft_model=draft_model,
        gamma=gamma,
    )

get_generation_backend

get_generation_backend(name='eager', *, engine=None, **kwargs)

Resolve a generation backend by registry name.

Backend-specific kwargs are forwarded to the factory — e.g. target_model=..., draft_model=..., gamma=... for the speculative backend, or engine=... for batched.

源代码位于： src/llm/generation/registry.py

def get_generation_backend(
    name: str = "eager",
    *,
    engine: ContinuousBatchingEngine | None = None,
    **kwargs: Any,
) -> GenerationBackend:
    """Resolve a generation backend by registry name.

    Backend-specific kwargs are forwarded to the factory — e.g.
    ``target_model=...``, ``draft_model=...``, ``gamma=...`` for the
    ``speculative`` backend, or ``engine=...`` for ``batched``.
    """
    ensure_backends_registered()
    return BACKEND_REGISTRY.get(name)(engine=engine, **kwargs)

Eager (Streaming) Backend

eager

stream_generate

stream_generate(model, tokenizer, prompt, max_new_tokens, temperature=1.0, top_k=None, top_p=None, repetition_penalty=1.0, frequency_penalty=0.0, presence_penalty=0.0, logit_bias=None, use_cache=True, stop=None)

Generator function for incremental text generation.

参数：

名称	类型	描述	默认
stop	str | list[str] | None	OpenAI-compat stop sequence(s). Generation halts the moment the accumulated output contains any of these as a suffix; the stop string itself is NOT included in the yielded output. Accepts a single string or a list of strings (OpenAI caps at 4). None is a no-op.	None

产生：

名称	类型	描述
str	Generator[str]	Newly generated text chunk (usually one token decoded).

源代码位于： src/llm/generation/eager.py

@torch.no_grad()
def stream_generate(
    model: DecoderModel,
    tokenizer: SimpleCharacterTokenizer,
    prompt: str,
    max_new_tokens: int,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    repetition_penalty: float = 1.0,
    frequency_penalty: float = 0.0,
    presence_penalty: float = 0.0,
    logit_bias: dict[int, float] | None = None,
    use_cache: bool = True,
    stop: str | list[str] | None = None,
) -> Generator[str]:
    """
    Generator function for incremental text generation.

    Args:
        stop: OpenAI-compat stop sequence(s). Generation halts the
            moment the accumulated output contains any of these as a
            suffix; the stop string itself is NOT included in the
            yielded output. Accepts a single string or a list of
            strings (OpenAI caps at 4). ``None`` is a no-op.

    yields:
        str: Newly generated text chunk (usually one token decoded).
    """
    model.eval()
    device = next(model.parameters()).device
    _reject_impossible_context(getattr(model, "max_seq_len", None), max_new_tokens)
    input_ids = tokenizer.encode(prompt)
    if not input_ids:
        # A prompt that decodes to zero tokens (empty string, or a tokenizer
        # with no char for it) would reach ``logits[0, -1, :]`` with a
        # zero-length sequence and crash with an opaque IndexError — and a
        # batch row never did the decode at all (round-71 empty-prompt fix).
        # The serving tier rejects an empty prompt up front; reject here too
        # for direct API callers.
        raise ValueError("prompt must decode to at least one token (got an empty token sequence)")
    input_tensor = torch.tensor(input_ids, dtype=torch.long, device=device).unsqueeze(0)
    max_seq_len = getattr(model, "max_seq_len", 512)
    kv_caches = create_decoder_kv_caches(model, batch_size=1) if use_cache else None

    # Prefill: truncate if needed to fit max_seq_len. Defensive: never slice
    # to an *empty* prompt — when max_new_tokens >= max_seq_len the slice
    # bound goes non-positive and the tensor becomes 0-length, crashing the
    # forward with a 500. Clamp to the last token so the model always sees a
    # non-empty context (the serving tier rejects this config up front).
    if input_tensor.size(1) + max_new_tokens > max_seq_len:
        keep = max(1, max_seq_len - max_new_tokens)
        input_tensor = input_tensor[:, -keep:]
        # Update input_ids to match truncated tensor
        input_ids = input_tensor[0].tolist()

    if use_cache:
        logits, kv_caches = model(input_tensor, kv_caches=kv_caches, use_cache=True)
        next_token_logits = logits[0, -1, :]
    else:
        # Initial forward pass without cache
        logits = model(input_tensor, use_cache=False)
        next_token_logits = logits[0, -1, :]

    _mask_pad_logits(next_token_logits, getattr(tokenizer, "pad_token_id", None))
    mask_undecodable_logits(next_token_logits, getattr(tokenizer, "vocab_size", None))

    generated_ids = input_ids.copy()

    # Stop-sequence tracking. We use a small buffer (``buffer``) that
    # holds decoded text not yet yielded to the caller. After each new
    # token is decoded we append it to the buffer and check whether the
    # buffer *ends with* any stop string (OpenAI semantics: generation
    # halts when a stop sequence appears as a suffix; the stop string
    # itself is NOT included in the output). If no stop is found, we
    # yield the portion of the buffer that extends beyond
    # ``max_stop_len`` characters from the end — that prefix is safe
    # because no stop sequence of length <= max_stop_len can span the
    # boundary. Only the last ``max_stop_len`` characters are kept
    # buffered so memory stays O(max_stop_len) regardless of how long
    # generation runs.
    stops = _normalize_stop(stop)
    max_stop_len = max((len(s) for s in stops), default=0) if stops else 0
    buffer = ""
    eos_id = getattr(tokenizer, "eos_token_id", None)

    for _ in range(max_new_tokens):
        if repetition_penalty != 1.0:
            next_token_logits = apply_repetition_penalty(next_token_logits, generated_ids, repetition_penalty)
        if frequency_penalty != 0.0:
            next_token_logits = apply_frequency_penalty(next_token_logits, generated_ids, frequency_penalty)
        if presence_penalty != 0.0:
            next_token_logits = apply_presence_penalty(next_token_logits, generated_ids, presence_penalty)
        if logit_bias:
            next_token_logits = apply_logit_bias(next_token_logits, logit_bias)

        token_id = sample_next_token(
            next_token_logits,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
        )
        # Model end-of-sequence: flush any buffered stop-prefix text and
        # halt. The EOS token itself is NOT part of the output (matches the
        # speculative backend's halting and standard LLM serving semantics);
        # without this the eager loop kept decoding through max_new_tokens
        # past EOS, emitting junk.
        if eos_id is not None and token_id == eos_id:
            if stops and buffer:
                yield buffer
            return
        generated_ids.append(token_id)
        text_chunk = tokenizer.decode([token_id])

        if stops and text_chunk:
            buffer += text_chunk
            # Check for a stop suffix — the first match wins.
            for s in stops:
                if buffer.endswith(s):
                    prefix = buffer[: len(buffer) - len(s)]
                    if prefix:
                        yield prefix
                    return
            # No stop found. Yield the safe prefix (everything beyond
            # the last max_stop_len characters) and keep the tail.
            if len(buffer) > max_stop_len:
                safe_len = len(buffer) - max_stop_len
                yield buffer[:safe_len]
                buffer = buffer[safe_len:]
        else:
            yield text_chunk

        next_input = torch.tensor([token_id], dtype=torch.long, device=device).unsqueeze(0)

        if use_cache:
            logits, kv_caches = model(next_input, kv_caches=kv_caches, use_cache=True)
            next_token_logits = logits[0, -1, :]
        else:
            # Without cache, append new token to full sequence and forward pass
            # generated_ids already has the new token appended
            full_input = torch.tensor(generated_ids, dtype=torch.long, device=device).unsqueeze(0)
            logits = model(full_input, use_cache=False)
            next_token_logits = logits[0, -1, :]

        _mask_pad_logits(next_token_logits, getattr(tokenizer, "pad_token_id", None))
        mask_undecodable_logits(next_token_logits, getattr(tokenizer, "vocab_size", None))

    # Flush any remaining buffered text after the loop ends (e.g. when
    # the buffer never exceeded max_stop_len or no stop sequence was found).
    if stops and buffer:
        yield buffer

generate

generate(model, tokenizer, prompt, max_new_tokens, temperature=1.0, top_k=None, top_p=None, repetition_penalty=1.0, frequency_penalty=0.0, presence_penalty=0.0, logit_bias=None, use_cache=True, stop=None)

Generate text from a prompt using a trained model.

源代码位于： src/llm/generation/eager.py

def generate(
    model: DecoderModel,
    tokenizer: SimpleCharacterTokenizer,
    prompt: str,
    max_new_tokens: int,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    repetition_penalty: float = 1.0,
    frequency_penalty: float = 0.0,
    presence_penalty: float = 0.0,
    logit_bias: dict[int, float] | None = None,
    use_cache: bool = True,
    stop: str | list[str] | None = None,
) -> str:
    """
    Generate text from a prompt using a trained model.
    """
    generator = stream_generate(
        model=model,
        tokenizer=tokenizer,
        prompt=prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        frequency_penalty=frequency_penalty,
        presence_penalty=presence_penalty,
        logit_bias=logit_bias,
        use_cache=use_cache,
        stop=stop,
    )
    return prompt + "".join(list(generator))

batch_generate

batch_generate(model, tokenizer, prompts, max_new_tokens, temperature=1.0, top_k=None, top_p=None, repetition_penalty=1.0, frequency_penalty=0.0, presence_penalty=0.0, logit_bias=None, stop=None)

Batch generate text from multiple prompts.

参数：

名称	类型	描述	默认
model	DecoderModel	The decoder model.	必需
tokenizer	SimpleCharacterTokenizer	The tokenizer.	必需
prompts	list[str]	List of input prompts.	必需
max_new_tokens	int	Maximum tokens to generate per prompt.	必需
temperature	float	Sampling temperature. 0 for greedy.	1.0
top_k	int | None	Top-k sampling parameter.	None
top_p	float | None	Nucleus sampling parameter.	None
repetition_penalty	float	Repetition penalty.	1.0
frequency_penalty	float	OpenAI-compatible per-frequency penalty (subtracts frequency_penalty * count(token) from each seen token's logit). 0.0 is a no-op.	0.0
presence_penalty	float	OpenAI-compatible per-presence penalty (subtracts a flat presence_penalty from each seen token's logit regardless of count). 0.0 is a no-op.	0.0
logit_bias	dict[int, float] | None	OpenAI-compatible additive per-token biases ({token_id: bias} added to the affected logits before sampling). None is a no-op.	None
stop	str | list[str] | None	OpenAI-compat stop sequence(s). Generation for each sequence halts the moment the generated text (post-prompt) contains any stop string; the stop string itself is NOT included in the returned text. Accepts a single string or a list of strings. None is a no-op.	None

返回：

类型	描述
list[str]	List of generated texts (prompt + generated tokens, with any
list[str]	stop sequence truncated).

源代码位于： src/llm/generation/eager.py

@torch.no_grad()
def batch_generate(
    model: DecoderModel,
    tokenizer: SimpleCharacterTokenizer,
    prompts: list[str],
    max_new_tokens: int,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    repetition_penalty: float = 1.0,
    frequency_penalty: float = 0.0,
    presence_penalty: float = 0.0,
    logit_bias: dict[int, float] | None = None,
    stop: str | list[str] | None = None,
) -> list[str]:
    """
    Batch generate text from multiple prompts.

    Args:
        model: The decoder model.
        tokenizer: The tokenizer.
        prompts: List of input prompts.
        max_new_tokens: Maximum tokens to generate per prompt.
        temperature: Sampling temperature. 0 for greedy.
        top_k: Top-k sampling parameter.
        top_p: Nucleus sampling parameter.
        repetition_penalty: Repetition penalty.
        frequency_penalty: OpenAI-compatible per-frequency penalty
            (subtracts ``frequency_penalty * count(token)`` from each
            seen token's logit). ``0.0`` is a no-op.
        presence_penalty: OpenAI-compatible per-presence penalty
            (subtracts a flat ``presence_penalty`` from each seen
            token's logit regardless of count). ``0.0`` is a no-op.
        logit_bias: OpenAI-compatible additive per-token biases
            (``{token_id: bias}`` added to the affected logits
            before sampling). ``None`` is a no-op.
        stop: OpenAI-compat stop sequence(s). Generation for each
            sequence halts the moment the generated text (post-prompt)
            contains any stop string; the stop string itself is NOT
            included in the returned text. Accepts a single string or
            a list of strings. ``None`` is a no-op.

    Returns:
        List of generated texts (prompt + generated tokens, with any
        stop sequence truncated).
    """
    if not prompts:
        return []

    model.eval()
    device = next(model.parameters()).device
    _reject_impossible_context(getattr(model, "max_seq_len", None), max_new_tokens)
    batch_size = len(prompts)

    # Encode all prompts. An empty row decodes to zero tokens and would either
    # crash the padded prefill or emit garbage (all-PAD prefill → arbitrary
    # sampled token, round-71 empty-prompt fix) — reject it up front.
    encoded_prompts = [tokenizer.encode(p) for p in prompts]
    for i, ids in enumerate(encoded_prompts):
        if not ids:
            raise ValueError(f"prompt[{i}] must decode to at least one token (got an empty token sequence)")

    # Truncate prompts that exceed ``max_seq_len - max_new_tokens`` **before**
    # padding and ``generated_ids`` initialisation.  Doing the truncate here
    # (instead of slicing the padded tensor afterwards) keeps
    # ``generated_ids`` in sync with the tokens the model actually attends
    # to in the prefill forward pass.  Otherwise the repetition-penalty
    # context would include token ids the model never saw.
    max_seq_len = getattr(model, "max_seq_len", 512)
    truncate_len = max_seq_len - max_new_tokens
    if truncate_len > 0:
        max_prompt_len = max(len(ids) for ids in encoded_prompts)
        if max_prompt_len + max_new_tokens > max_seq_len:
            encoded_prompts = [ids[-truncate_len:] if len(ids) > truncate_len else ids for ids in encoded_prompts]

    prompt_lengths = [len(p) for p in encoded_prompts]
    max_prompt_len = max(prompt_lengths) if prompt_lengths else 0

    # Get pad token id
    pad_id = getattr(tokenizer, "pad_token_id", 0)

    # Left-pad sequences to align generation positions
    padded_inputs = []
    for ids in encoded_prompts:
        padding_len = max_prompt_len - len(ids)
        padded_inputs.append([pad_id] * padding_len + ids)

    input_tensor = torch.tensor(padded_inputs, dtype=torch.long, device=device)

    # Left-pad attention mask (True = mask out, matching the codebase SDPA
    # convention in ``llm.core.attn.sdpa``).  The left-pad columns are real
    # pad-token embeddings under the (default) causal mask: without an
    # explicit mask the prefill forward attends over the pad K/V, AND those
    # pad columns stay in the KV cache for every decode step — silently
    # diverging from the single-prompt path (RIL ISS-070).  We build one
    # mask sized to the full generation window and slice it per forward:
    # the prefill key length is ``max_prompt_len`` and each decode step t
    # grows the key length to ``max_prompt_len + t + 1``.  Generated columns
    # (beyond ``max_prompt_len``) are never masked.  Slices stay 4-D
    # ``[B, 1, 1, k_len]`` so they broadcast to ``[B, N, Sq, Sk]`` like the
    # batch-engine's ``run_attn_mask``.
    max_total_len = max_prompt_len + max_new_tokens
    pad_mask = torch.zeros((batch_size, 1, 1, max_total_len), dtype=torch.bool, device=device)
    for i, ids in enumerate(encoded_prompts):
        pad_len = max_prompt_len - len(ids)
        if pad_len > 0:
            pad_mask[i, 0, 0, :pad_len] = True

    # Track generated ids per sequence — seeded from the (possibly truncated)
    # encoded prompts so the repetition-penalty context matches the model's
    # actual prefill input.
    generated_ids: list[list[int]] = [ids.copy() for ids in encoded_prompts]

    kv_caches = create_decoder_kv_caches(model, batch_size=batch_size)
    logits, kv_caches = model(
        input_tensor,
        kv_caches=kv_caches,
        use_cache=True,
        attn_mask=pad_mask[..., :max_prompt_len],
    )
    next_token_logits = logits[:, -1, :]  # [B, vocab_size]

    _mask_pad_logits(next_token_logits, getattr(tokenizer, "pad_token_id", None))
    mask_undecodable_logits(next_token_logits, getattr(tokenizer, "vocab_size", None))

    for step in range(max_new_tokens):
        for i in range(batch_size):
            row_logits = next_token_logits[i]
            if repetition_penalty != 1.0:
                row_logits = apply_repetition_penalty(row_logits, generated_ids[i], repetition_penalty)
            if frequency_penalty != 0.0:
                row_logits = apply_frequency_penalty(row_logits, generated_ids[i], frequency_penalty)
            if presence_penalty != 0.0:
                row_logits = apply_presence_penalty(row_logits, generated_ids[i], presence_penalty)
            if logit_bias:
                row_logits = apply_logit_bias(row_logits, logit_bias)
            token_id = sample_next_token(
                row_logits,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
            )
            generated_ids[i].append(token_id)

        next_tokens = torch.tensor(
            [[generated_ids[i][-1]] for i in range(batch_size)],
            dtype=torch.long,
            device=device,
        )

        # Decode key length grows by one per step (the cache already holds
        # ``max_prompt_len + step`` keys after the prefill, plus the new one).
        logits, kv_caches = model(
            next_tokens,
            kv_caches=kv_caches,
            use_cache=True,
            attn_mask=pad_mask[..., : max_prompt_len + step + 1],
        )
        next_token_logits = logits[:, -1, :]

        _mask_pad_logits(next_token_logits, getattr(tokenizer, "pad_token_id", None))
        mask_undecodable_logits(next_token_logits, getattr(tokenizer, "vocab_size", None))

    # Truncate each sequence at its first EOS so both decode paths below
    # omit the EOS token and any junk generated after it (a sequence that
    # already finished keeps occupying its batch slot, but its tail is cut
    # here). Matches stream_generate / the speculative backend.
    eos_id = getattr(tokenizer, "eos_token_id", None)
    if eos_id is not None:
        for i in range(batch_size):
            gen_start = len(encoded_prompts[i])
            for j in range(gen_start, len(generated_ids[i])):
                if generated_ids[i][j] == eos_id:
                    del generated_ids[i][j:]
                    break

    # Decode results, applying stop sequences when provided.
    # OpenAI semantics: generation halts when a stop sequence appears as
    # a **suffix** of the running output. We simulate incremental decode
    # to find the first suffix match — .find() would match anywhere and
    # could prematurely truncate on prompt-embedded sequences or matches
    # that wouldn't have been a suffix during streaming.
    stops = _normalize_stop(stop)
    if stops:
        prompt_texts = [tokenizer.decode(p) for p in encoded_prompts]
        results = []
        for i in range(batch_size):
            running = prompt_texts[i]
            p_len = len(prompt_texts[i])
            generated_part = ""
            # Walk generated tokens one by one, checking for suffix stop
            # after each decode (mirrors stream_generate incremental logic).
            gen_start = len(encoded_prompts[i])
            for tid in generated_ids[i][gen_start:]:
                running += tokenizer.decode([tid])
                generated_part = running[p_len:]
                truncated = False
                for s in stops:
                    if generated_part.endswith(s):
                        generated_part = generated_part[: -len(s)]
                        truncated = True
                        break
                if truncated:
                    break
            results.append(prompt_texts[i] + generated_part)
        return results

    return [tokenizer.decode(ids) for ids in generated_ids]

Speculative Decoding Backend

speculative

Speculative decoding (Leviathan et al., 2023).

A small draft model speculates gamma candidate tokens ahead of the target model. The target then scores all gamma + 1 positions in a single forward pass, and the algorithm either accepts each candidate (with probability preserving the target distribution) or samples a correction token. Net effect: every accepted token costs roughly one draft forward; only rejections require the more expensive target forward.

The implementation is greedy/sample-aware via :func:llm.generation.sampling.sample_next_token and emits decoded chunks through the standard generator protocol so it slots into the existing :class:llm.generation.backends.GenerationBackend.

References

Leviathan, Kalman, Matan Kalman, and Yossi Matias. "Fast Inference from Transformers via Speculative Decoding." ICML 2023. https://arxiv.org/abs/2211.17192

TokenizerLike

Bases: Protocol

Anything with encode/decode + optional pad/eos token ids.

源代码位于： src/llm/generation/speculative.py

class TokenizerLike(Protocol):
    """Anything with encode/decode + optional pad/eos token ids."""

    eos_token_id: int | None
    pad_token_id: int | None

    def encode(self, text: str, /) -> list[int]: ...
    def decode(self, tokens: list[int], /) -> str: ...

speculative_generate

speculative_generate(target, draft, tokenizer, prompt, max_new_tokens, *, gamma=5, temperature=1.0, top_k=None, top_p=None, repetition_penalty=1.0, frequency_penalty=0.0, presence_penalty=0.0, logit_bias=None, seed=None, stop=None)

Speculative decoding generator.

Yields decoded chunks. Stops after max_new_tokens produced tokens or on EOS.

参数：

名称	类型	描述	默认
target	DecoderModel	Target model (the "expensive" one). Its forward distribution is the canonical output distribution.	必需
draft	DecoderModel	Draft model (the "cheap" one). Must share vocabulary with the target and have the same max_seq_len (or larger - we only enforce the prompt fits).	必需
tokenizer	TokenizerLike	Tokenizer with encode, decode, pad_token_id, eos_token_id.	必需
prompt	str	Prompt text.	必需
max_new_tokens	int	Hard cap on generated tokens.	必需
gamma	int	Number of speculative tokens per round. Typical values are 4-8.	5
temperature	float	Sampling temperature for the correction token (the algorithm preserves the target distribution under these settings).	1.0
top_k	int | None	Top-k sampling parameter for the correction token.	None
top_p	float | None	Nucleus-sampling (top-p) parameter for the correction token.	None
repetition_penalty	float	Applied to both draft and target logits before sampling.	1.0
seed	int | None	Optional RNG seed for reproducible rejection sampling.	None
stop	str | list[str] | None	OpenAI-compat stop sequence(s). Generation halts the moment the accumulated output contains any of these as a suffix; the stop string itself is NOT included in the yielded output. Accepts a single string or a list of strings. None is a no-op.	None

源代码位于： src/llm/generation/speculative.py

@torch.no_grad()
def speculative_generate(
    target: DecoderModel,
    draft: DecoderModel,
    tokenizer: TokenizerLike,
    prompt: str,
    max_new_tokens: int,
    *,
    gamma: int = 5,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    repetition_penalty: float = 1.0,
    frequency_penalty: float = 0.0,
    presence_penalty: float = 0.0,
    logit_bias: dict[int, float] | None = None,
    seed: int | None = None,
    stop: str | list[str] | None = None,
):
    """Speculative decoding generator.

    Yields decoded chunks. Stops after ``max_new_tokens`` produced
    tokens or on EOS.

    Args:
        target: Target model (the "expensive" one). Its forward
            distribution is the canonical output distribution.
        draft: Draft model (the "cheap" one). Must share vocabulary
            with the target and have the same ``max_seq_len`` (or
            larger - we only enforce the prompt fits).
        tokenizer: Tokenizer with ``encode``, ``decode``,
            ``pad_token_id``, ``eos_token_id``.
        prompt: Prompt text.
        max_new_tokens: Hard cap on generated tokens.
        gamma: Number of speculative tokens per round. Typical
            values are 4-8.
        temperature: Sampling temperature for the **correction**
            token (the algorithm preserves the target distribution
            under these settings).
        top_k: Top-k sampling parameter for the correction token.
        top_p: Nucleus-sampling (top-p) parameter for the correction
            token.
        repetition_penalty: Applied to both draft and target logits
            before sampling.
        seed: Optional RNG seed for reproducible rejection sampling.
        stop: OpenAI-compat stop sequence(s). Generation halts the
            moment the accumulated output contains any of these as a
            suffix; the stop string itself is NOT included in the
            yielded output. Accepts a single string or a list of
            strings. ``None`` is a no-op.
    """
    if gamma < 1:
        raise ValueError(f"gamma must be >= 1, got {gamma}")
    if seed is not None:
        torch.manual_seed(seed)

    target.eval()
    draft.eval()
    device = next(target.parameters()).device

    # Same context-window handshake as the eager backend (RIL ISS-124): a
    # budget that cannot fit is rejected up front, and an over-long prompt is
    # truncated to fit. The speculative backend keeps rebuilding the full
    # context tensor each round and forwarding it through the learned
    # positional-encoding table, so without this a long prompt crashed with
    # ``ValueError: Sequence endpoint N exceeds maximum sequence length``
    # while the eager backend served the same input fine (it truncates).
    max_seq_len = getattr(target, "max_seq_len", None)
    _reject_impossible_context(max_seq_len, max_new_tokens)
    prompt_ids = tokenizer.encode(prompt)
    if max_seq_len is not None and len(prompt_ids) + max_new_tokens > max_seq_len:
        keep = max(1, max_seq_len - max_new_tokens)
        prompt_ids = prompt_ids[-keep:]

    generated_ids: list[int] = list(prompt_ids)
    eos_id = getattr(tokenizer, "eos_token_id", None)

    # Stop-sequence tracking via a small suffix buffer (same strategy
    # as stream_generate: keep at most ``max_stop_len`` chars un-yielded
    # so the buffer is O(max_stop_len) and suffix matching is exact).
    stops = _normalize_stop(stop)
    max_stop_len = max((len(s) for s in stops), default=0) if stops else 0
    buffer = ""

    while len(generated_ids) - len(prompt_ids) < max_new_tokens:
        # 1. Draft: generate gamma candidates with the small model.
        # We rebuild the context tensor at each step so the draft
        # can use its KV cache naturally.
        draft_ids = list(generated_ids)
        draft_tokens: list[int] = []
        draft.eval()
        for _ in range(gamma):
            ctx = torch.tensor([draft_ids], dtype=torch.long, device=device)
            draft_out = draft(ctx, use_cache=False)
            logits = draft_out[0] if isinstance(draft_out, tuple) else draft_out
            next_logits = logits[0, -1, :]
            mask_undecodable_logits(next_logits, getattr(tokenizer, "vocab_size", None))
            # Mask the PAD sentinel so the draft never proposes it as a
            # candidate (round-71 speculative fix; eager masks every step).
            _mask_pad_logits(next_logits, getattr(tokenizer, "pad_token_id", None))
            if repetition_penalty != 1.0:
                next_logits = apply_repetition_penalty(next_logits, draft_ids, repetition_penalty)
            if frequency_penalty != 0.0:
                next_logits = apply_frequency_penalty(next_logits, draft_ids, frequency_penalty)
            if presence_penalty != 0.0:
                next_logits = apply_presence_penalty(next_logits, draft_ids, presence_penalty)
            if logit_bias:
                next_logits = apply_logit_bias(next_logits, logit_bias)
            tok = sample_next_token(
                next_logits,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
            )
            draft_tokens.append(tok)
            draft_ids.append(tok)
            if eos_id is not None and tok == eos_id:
                break

        # 2. Verify against the target.
        accept_count, bonus = _verify_speculative_tokens(
            target,
            draft,
            torch.tensor([generated_ids], dtype=torch.long, device=device),
            draft_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            frequency_penalty=frequency_penalty,
            presence_penalty=presence_penalty,
            logit_bias=logit_bias,
            tokenizer_vocab_size=getattr(tokenizer, "vocab_size", None),
            pad_token_id=getattr(tokenizer, "pad_token_id", None),
        )

        # 3. Emit accepted tokens + bonus (or correction). The EOS token is
        # never part of the yielded output — halt on it *before* decoding and
        # appending, matching the eager backend (RIL ISS-96/ISS-98) which
        # stops on EOS without emitting the EOS token's decoded text.
        for i in range(accept_count):
            tok = draft_tokens[i]
            if eos_id is not None and tok == eos_id:
                if stops and buffer:
                    yield buffer
                return
            generated_ids.append(tok)
            text_chunk = tokenizer.decode([tok])
            if stops and text_chunk:
                buffer += text_chunk
                for s in stops:
                    if buffer.endswith(s):
                        prefix = buffer[: len(buffer) - len(s)]
                        if prefix:
                            yield prefix
                        return
                if len(buffer) > max_stop_len:
                    safe_len = len(buffer) - max_stop_len
                    yield buffer[:safe_len]
                    buffer = buffer[safe_len:]
            else:
                yield text_chunk
            if len(generated_ids) - len(prompt_ids) >= max_new_tokens:
                if stops and buffer:
                    yield buffer
                return

        # Append the bonus or correction token (one per round).
        if bonus is not None:
            if eos_id is not None and bonus == eos_id:
                if stops and buffer:
                    yield buffer
                return
            generated_ids.append(bonus)
            text_chunk = tokenizer.decode([bonus])
            if stops and text_chunk:
                buffer += text_chunk
                for s in stops:
                    if buffer.endswith(s):
                        prefix = buffer[: len(buffer) - len(s)]
                        if prefix:
                            yield prefix
                        return
                if len(buffer) > max_stop_len:
                    safe_len = len(buffer) - max_stop_len
                    yield buffer[:safe_len]
                    buffer = buffer[safe_len:]
            else:
                yield text_chunk

    # Flush any remaining buffered text when the loop exhausts
    # max_new_tokens without a stop or EOS triggering an early return.
    if stops and buffer:
        yield buffer
    return