跳转至

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.

源代码位于： src/llm/generation/sampling.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def apply_repetition_penalty(
    logits: torch.Tensor,
    token_ids: list[int],
    repetition_penalty: float,
) -> torch.Tensor:
    """Apply repetition penalty in-place on 1D logits."""
    if repetition_penalty == 1.0 or not token_ids:
        return logits

    adjusted = logits.clone()
    device = adjusted.device
    ids = torch.tensor(token_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 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
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 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
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

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 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
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())

    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")

    probs = torch.softmax(next_logits, dim=-1)
    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 
16
17
18
19
20
21
22
23
24
25
26
27
28
@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

GenerationBackend

Bases: ABC

Backend protocol for text generation.

源代码位于： src/llm/generation/backends.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class GenerationBackend(abc.ABC):
    """Backend protocol for text generation."""

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

    def generate(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
    ) -> str:
        chunks = list(self.stream(model, tokenizer, prompt, config))
        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 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class EagerGenerationBackend(GenerationBackend):
    """Default in-process generation using the library stream_generate path."""

    def stream(
        self,
        model: DecoderModel,
        tokenizer: Any,
        prompt: str,
        config: GenerationConfig,
    ) -> 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,
        )

    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=config.logit_bias,
        )

BatchedGenerationBackend

Bases: GenerationBackend

Generation via ContinuousBatchingEngine (iteration-level scheduling).

源代码位于： src/llm/generation/backends.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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,
    ) -> Generator[str]:
        from llm.serving.schemas import GenerationRequest

        request = 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,
        )
        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,
            )
            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 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
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,
    ) -> 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=config.logit_bias,
        )

    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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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)

Generator function for incremental text generation.

产生：

名称 类型 描述
str Generator[str]

Newly generated text chunk (usually one token decoded).
源代码位于： src/llm/generation/eager.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@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,
) -> Generator[str]:
    """
    Generator function for incremental text generation.

    yields:
        str: Newly generated text chunk (usually one token decoded).
    """
    model.eval()
    device = next(model.parameters()).device
    input_ids = tokenizer.encode(prompt)
    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
    if input_tensor.size(1) + max_new_tokens > max_seq_len:
        input_tensor = input_tensor[:, -(max_seq_len - max_new_tokens) :]
        # 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))

    generated_ids = input_ids.copy()

    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,
        )
        generated_ids.append(token_id)
        text_chunk = tokenizer.decode([token_id])
        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))

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)

Generate text from a prompt using a trained model.

源代码位于： src/llm/generation/eager.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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,
) -> 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,
    )
    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)

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

返回：

类型 描述
list[str]

List of generated texts (prompt + generated tokens).
源代码位于： src/llm/generation/eager.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
@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,
) -> 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.

    Returns:
        List of generated texts (prompt + generated tokens).
    """
    if not prompts:
        return []

    model.eval()
    device = next(model.parameters()).device
    batch_size = len(prompts)

    # Encode all prompts
    encoded_prompts = [tokenizer.encode(p) for p in prompts]
    prompt_lengths = [len(p) for p in encoded_prompts]
    max_prompt_len = max(prompt_lengths)

    # 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)

    # Track generated ids per sequence
    generated_ids: list[list[int]] = [ids.copy() for ids in encoded_prompts]

    # Prefill
    max_seq_len = getattr(model, "max_seq_len", 512)
    if input_tensor.size(1) + max_new_tokens > max_seq_len:
        truncate_len = max_seq_len - max_new_tokens
        input_tensor = input_tensor[:, -truncate_len:]

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

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

    for _ 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,
        )

        logits, kv_caches = model(next_tokens, kv_caches=kv_caches, use_cache=True)
        next_token_logits = logits[:, -1, :]

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

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