| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This example implements the DRY (Don't Repeat Yourself) repetition penalty as |
| 4 | a Model Runner V2 custom logits processor, ``DryState``. |
| 5 | |
| 6 | A token that would extend a repeat of at least ``allowed_length`` tokens loses |
| 7 | ``multiplier * base ** (repeat_length - allowed_length)`` from its logit, with the |
| 8 | exponent clamped as in llama.cpp. |
| 9 | Parameter names and matching semantics follow llama.cpp's |
| 10 | ``llama_sampler_dry_apply``. llama.cpp's version is itself ported from pi6am's |
| 11 | koboldcpp implementation of p-e-w's scheme for text-generation-webui. |
| 12 | |
| 13 | Pass the class to ``LLM`` and configure it per request through ``extra_args``:: |
| 14 | |
| 15 | LLM(..., logits_processors=[DryState]) |
| 16 | SamplingParams(extra_args={"dry_multiplier": 0.8}) |
| 17 | |
| 18 | The recognized keys are ``dry_multiplier`` (0.0, off), ``dry_base`` (1.75), |
| 19 | ``dry_allowed_length`` (2), ``dry_penalty_last_n`` (-1, the whole context) |
| 20 | and ``dry_sequence_breakers``. This processor claims the ``dry_*`` namespace, |
| 21 | so any other ``dry_*`` key fails the request. Speculative decoding is not |
| 22 | supported and is refused at construction. |
| 23 | |
| 24 | Run the example with:: |
| 25 | |
| 26 | python examples/features/logits_processor/dry.py |
| 27 | |
| 28 | It sends one prompt twice in one batch, once without DRY and once with it, and |
| 29 | prints both outputs, yielding an output similar to that shown below: |
| 30 | |
| 31 | Generated Outputs: |
| 32 | ------------------------------------------------------------ |
| 33 | Prompt: 'The future of AI is', without DRY |
| 34 | Output: ' in the hands of the people.\n\nThe future of AI is in the hands of |
| 35 | the people.\n\nThe future of AI is in the hands of the |
| 36 | people.\n\nThe future of AI is in the hands of the people.\n\nThe |
| 37 | future of AI is in the hands of the people.\n' |
| 38 | ------------------------------------------------------------ |
| 39 | Prompt: 'The future of AI is', with DRY |
| 40 | Output: ' in the hands of the people.\n\nThe future of AI is the future of |
| 41 | the human race.\n\nThe future of AI is a future of the human race, |
| 42 | and the future of humanity.\n\nThe future of AI is an AI that is |
| 43 | capable of making decisions that are not based on human |
| 44 | intelligence.' |
| 45 | ------------------------------------------------------------ |
| 46 | """ |
| 47 | |
| 48 | import math |
| 49 | import weakref |
| 50 | from typing import TYPE_CHECKING, Any, NamedTuple |
| 51 | |
| 52 | import numpy as np |
| 53 | import torch |
| 54 | |
| 55 | from vllm.logger import init_logger |
| 56 | from vllm.sampling_params import SamplingParams |
| 57 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 58 | from vllm.utils.torch_utils import async_tensor_h2d |
| 59 | from vllm.v1.worker.gpu.sample.logits_processor import ( |
| 60 | LogitsContext, |
| 61 | LogitsProcessor, |
| 62 | LogitsProcRequestState, |
| 63 | ) |
| 64 | |
| 65 | if TYPE_CHECKING: |
| 66 | from vllm.config import VllmConfig |
| 67 | |
| 68 | logger = init_logger(__name__) |
| 69 | |
| 70 | _J_BUDGET = 2048 |
| 71 | # Run lengths feed an int16 sum, so they must stay below 2**15. |
| 72 | assert _J_BUDGET < 2**15, "_J_BUDGET must fit the int16 run-length sum" |
| 73 | |
| 74 | # Byte budget for the per-chunk transients of the match scan. The gather |
| 75 | # W32[:, idx1] materializes [R, chunk, J] int32 (4 B/elem) before the |
| 76 | # comparison reduces it to bool; with the masks, the int8 cumprod and |
| 77 | # the int16 row sums, the marginal transient cost is ~8 B/elem. 24 keeps |
| 78 | # headroom for the fixed R*vocab penalty accumulator floor and allocator slack. |
| 79 | _CHUNK_BYTE_BUDGET = 256 * 1024 * 1024 |
| 80 | _CHUNK_PEAK_BYTES_PER_ELEM = 24 |
| 81 | |
| 82 | # llama.cpp's FLOAT_MAX_LOG (src/llama-sampler.cpp): ln(float32 max). |
| 83 | _FLOAT_MAX_LOG = 88.7228391 |
| 84 | |
| 85 | # Breakers are truncated to this many code points. llama.cpp cuts at this many |
| 86 | # bytes (llama_sampler_init_dry), which keeps fewer characters of a non-ASCII |
| 87 | # breaker. |
| 88 | _MAX_BREAKER_CHAR_LEN = 40 |
| 89 | |
| 90 | # Upper bound on distinct breaker sets cached per tokenizer. |
| 91 | _MAX_CACHED_BREAKER_SETS = 64 |
| 92 | |
| 93 | STR_SPEC_DEC_REJECTS_DRY = ( |
| 94 | "The DRY logits processor is not supported when speculative decoding is enabled." |
| 95 | ) |
| 96 | |
| 97 | DEFAULT_DRY_SEQUENCE_BREAKERS = ("\n", ":", '"', "*") |
| 98 | """llama.cpp's default DRY sequence breakers.""" |
| 99 | |
| 100 | MAX_DRY_SEQUENCE_BREAKERS = 64 |
| 101 | """Upper bound on the dry_sequence_breakers list length. Each multi-character |
| 102 | breaker in an uncached set scans the vocabulary.""" |
| 103 | |
| 104 | _DRY_INT_MAX = 2**31 - 1 |
| 105 | """Upper bound on the integral DRY parameters, matching llama-server's |
| 106 | INT32_MAX cap. The state below stores them in an int64 numpy array, where |
| 107 | anything at or above 2**63 raises OverflowError inside execute_model and |
| 108 | takes the engine down.""" |
| 109 | |
| 110 | _MAX_CACHED_BREAKER_MASKS = 64 |
| 111 | """Upper bound on distinct breaker sets holding a device mask, mirroring |
| 112 | _MAX_CACHED_BREAKER_SETS on the host side.""" |
| 113 | |
| 114 | _WARMUP_WINDOW = 8 |
| 115 | """Window length of the dry_core call in DryState.__init__, long enough to |
| 116 | hold a match at the default dry_allowed_length.""" |
| 117 | |
| 118 | _DRY_DEFAULTS: dict[str, Any] = { |
| 119 | "dry_multiplier": 0.0, |
| 120 | "dry_base": 1.75, |
| 121 | "dry_allowed_length": 2, |
| 122 | "dry_penalty_last_n": -1, # whole context; llama.cpp's default is 64 |
| 123 | "dry_sequence_breakers": DEFAULT_DRY_SEQUENCE_BREAKERS, |
| 124 | } |
| 125 | """Recognized extra_args keys and their defaults.""" |
| 126 | |
| 127 | |
| 128 | # --------------------------------------------------------------------------- |
| 129 | # The match computation |
| 130 | # --------------------------------------------------------------------------- |
| 131 | # ``_dry_penalties`` is a sequential port of llama.cpp's Z-algorithm, used as the |
| 132 | # reference. ``dry_core`` is the vectorized form: the comparisons that determine |
| 133 | # the match length ending at position ``i`` lie along one diagonal, so a |
| 134 | # cumprod-sum along ``j`` of an ``[R, K, J]`` tensor gives every match length |
| 135 | # without a sequential scan. Capping ``J`` at ``allowed_length + max_exponent`` |
| 136 | # is exact, because the exponent clamp maps every longer match to the same |
| 137 | # penalty. Requests whose cap is unusable (``max_exponent == 0``, i.e. |
| 138 | # ``base <= 1.000001``, or a cap beyond ``_J_BUDGET``) use the sequential form. |
| 139 | |
| 140 | |
| 141 | def max_exponent(base: float) -> int: |
| 142 | """Exponent clamp, mirroring llama.cpp bit-for-bit. |
| 143 | |
| 144 | llama.cpp computes ``FLOAT_MAX_LOG / std::log(dry_base)`` entirely in |
| 145 | float32. Computing the quotient in float64 lands on the wrong side of |
| 146 | the integer truncation for some bases: at ``base=2.0`` float64 gives |
| 147 | 127.99999998 -> 127 while float32 gives exactly 128.0 -> 128. |
| 148 | """ |
| 149 | if base <= 1.000001: |
| 150 | return 0 |
| 151 | return int(np.float32(_FLOAT_MAX_LOG) / np.log(np.float32(base))) |
| 152 | |
| 153 | |
| 154 | def _dry_penalties( |
| 155 | window: list[int], |
| 156 | breakers: frozenset[int], |
| 157 | multiplier: float, |
| 158 | base: float, |
| 159 | allowed_length: int, |
| 160 | max_exp: int, |
| 161 | ) -> dict[int, float]: |
| 162 | """Compute DRY penalties for one request. |
| 163 | |
| 164 | Port of the scan in ``llama_sampler_dry_apply``: a reverse-direction |
| 165 | Z-algorithm finds, for each position, the length of the match between |
| 166 | the window suffix and the sequence ending at that position; each |
| 167 | match's follower token is charged the penalty for the longest repeat |
| 168 | it would extend. |
| 169 | |
| 170 | Returns a dict mapping token id -> penalty to subtract from its logit. |
| 171 | """ |
| 172 | m = len(window) |
| 173 | |
| 174 | def rat(i: int) -> int: # reverse access: rat(0) == last token |
| 175 | return window[m - 1 - i] |
| 176 | |
| 177 | # Step 1: the nearest breaker from the end caps the match length. |
| 178 | rep_limit = m |
| 179 | for i in range(m): |
| 180 | if rat(i) in breakers: |
| 181 | rep_limit = i |
| 182 | break |
| 183 | if rep_limit < allowed_length: |
| 184 | return {} |
| 185 | |
| 186 | # Step 2: reverse-direction Z-algorithm -> per-position repeat counts |
| 187 | # (forward-indexed into ``window`` via cnt[last - k]). |
| 188 | cnt = [0] * m |
| 189 | last = m - 1 |
| 190 | lt = rt = 0 |
| 191 | for k in range(1, m): |
| 192 | if k > rt: |
| 193 | # Outside the current Z-box: extend naively. |
| 194 | z = 0 |
| 195 | while z + k < m and rat(z) == rat(z + k): |
| 196 | z += 1 |
| 197 | cnt[last - k] = min(z, rep_limit) |
| 198 | if z > 0: |
| 199 | lt, rt = k, k + z - 1 |
| 200 | else: |
| 201 | p = k - lt |
| 202 | right_part_len = rt - k + 1 |
| 203 | if cnt[last - p] < right_part_len: |
| 204 | # Fully inside the Z-box: copy. |
| 205 | cnt[last - k] = min(cnt[last - p], rep_limit) |
| 206 | else: |
| 207 | # Touches the right edge: extend past it. |
| 208 | j = rt + 1 |
| 209 | while j < m and rat(j) == rat(j - k): |
| 210 | j += 1 |
| 211 | cnt[last - k] = min(j - k, rep_limit) |
| 212 | lt, rt = k, j - 1 |
| 213 | |
| 214 | # Step 3: map each repeat's follower token to the longest repeat that |
| 215 | # it would extend. |
| 216 | max_token_repeat: dict[int, int] = {} |
| 217 | for i in range(m - 1): |
| 218 | repeat_len = cnt[i] |
| 219 | if repeat_len >= allowed_length: |
| 220 | tok = window[i + 1] |
| 221 | if max_token_repeat.get(tok, -1) < repeat_len: |
| 222 | max_token_repeat[tok] = repeat_len |
| 223 | |
| 224 | # Step 4: exponential penalty, exponent clamped for float32 safety. |
| 225 | # Breaker tokens are never penalized. A value overflowing float32 |
| 226 | # saturates the logit to -inf downstream, as in llama.cpp. |
| 227 | penalties: dict[int, float] = {} |
| 228 | for tok, repeat_len in max_token_repeat.items(): |
| 229 | if tok in breakers: |
| 230 | continue |
| 231 | exponent = repeat_len - allowed_length |
| 232 | if max_exp and exponent > max_exp: |
| 233 | exponent = max_exp |
| 234 | penalties[tok] = multiplier * (base**exponent) |
| 235 | return penalties |
| 236 | |
| 237 | |
| 238 | def dry_core( |
| 239 | logits: torch.Tensor, |
| 240 | row_idx: torch.Tensor, |
| 241 | W: torch.Tensor, |
| 242 | n_r: torch.Tensor, |
| 243 | allowed: torch.Tensor, |
| 244 | max_exp: torch.Tensor, |
| 245 | mult: torch.Tensor, |
| 246 | base: torch.Tensor, |
| 247 | breaker_masks: list[torch.Tensor | None], |
| 248 | j_budget: int, |
| 249 | ) -> torch.Tensor: |
| 250 | """Apply DRY penalties in place given batched window tensors. |
| 251 | |
| 252 | The tensor bounds below are not checked, to avoid a sync. |
| 253 | |
| 254 | Args: |
| 255 | logits: [B, vocab] float tensor, modified in place. |
| 256 | row_idx: [R] int64, row of ``logits`` for each DRY request. |
| 257 | W: [R, N] int64 windows, right-aligned (window tokens occupy the |
| 258 | trailing ``n_r`` columns; leading columns are ignored via ``n_r`` |
| 259 | masks, whatever they hold). |
| 260 | n_r: [R] int64 window lengths. |
| 261 | allowed: [R] int64 per-request allowed length. |
| 262 | max_exp: [R] int64 per-request exponent ceiling, > 0. |
| 263 | mult: [R] float32 per-request multiplier, >= 0, already rounded |
| 264 | through float32, as llama.cpp stores it. |
| 265 | base: [R] float32 per-request base, >= 1, rounded the same way. |
| 266 | breaker_masks: per-request [vocab] bool masks (or None for no |
| 267 | breakers). ``vocab`` must equal ``logits.shape[-1]``. |
| 268 | j_budget: max(allowed + max_exp), <= _J_BUDGET, computed by the caller on |
| 269 | the host. |
| 270 | |
| 271 | """ |
| 272 | device = logits.device |
| 273 | vocab = logits.shape[-1] |
| 274 | R, N = W.shape |
| 275 | |
| 276 | # j_budget arrives as an unchecked host int; this check fails loudly rather |
| 277 | # than wrapping an int16 run length. The routing predicate in DryState.apply |
| 278 | # is what actually bounds it. |
| 279 | if j_budget > _J_BUDGET: |
| 280 | raise ValueError(f"j_budget {j_budget} over _J_BUDGET {_J_BUDGET}") |
| 281 | # base >= 1 and mult >= 0 are preconditions too (see the amax comment below). |
| 282 | # They live in device tensors and reading them back would sync; use_dry() |
| 283 | # and validate_params enforce them instead. |
| 284 | |
| 285 | # rep_limit: distance from the end of the nearest breaker (llama.cpp step 1). |
| 286 | # A breaker at column c is j = N-1-c tokens from the end. |
| 287 | rep_limit = n_r.clone() |
| 288 | # bool(bm.any()) would be a per-request host sync; the None check says the same. |
| 289 | any_breakers = any(bm is not None for bm in breaker_masks) |
| 290 | bmask = None |
| 291 | if any_breakers: |
| 292 | # Stacked once, used twice: nearest-breaker search here, penalty zeroing |
| 293 | # after the scan. The masks are cached and resident (DryState._breaker_masks), |
| 294 | # so the stack costs one byte per (row, vocab) entry. |
| 295 | bmask = torch.stack( |
| 296 | [ |
| 297 | bm |
| 298 | if bm is not None |
| 299 | else torch.zeros(vocab, dtype=torch.bool, device=device) |
| 300 | for bm in breaker_masks |
| 301 | ] |
| 302 | ) |
| 303 | Bwin = bmask.gather(1, W.clamp(min=0)) |
| 304 | valid_cols = torch.arange(N, device=device)[None, :] >= (N - n_r)[:, None] |
| 305 | Bwin &= valid_cols |
| 306 | has_breaker = Bwin.any(dim=1) |
| 307 | # max breaker column -> nearest to the end. |
| 308 | max_col = torch.where( |
| 309 | has_breaker, |
| 310 | (Bwin * torch.arange(1, N + 1, device=device)[None, :]).max(dim=1).values |
| 311 | - 1, |
| 312 | torch.zeros_like(n_r), |
| 313 | ) |
| 314 | rep_limit = torch.where(has_breaker, (N - 1) - max_col, n_r) |
| 315 | |
| 316 | # llama.cpp: if rep_limit < allowed_length, the request produces nothing. |
| 317 | active = rep_limit >= allowed |
| 318 | |
| 319 | # Token ids fit int32, and gathering int32 halves the dominant per-chunk transient. |
| 320 | W32 = W.to(torch.int32) |
| 321 | # J is a tensor shape, so it must be known host-side; the caller computes it |
| 322 | # from numpy (see DryState.apply) to avoid a per-step device readback. |
| 323 | J = max(1, min(j_budget, N)) |
| 324 | idx2 = torch.arange(N - 1, N - 1 - J, -1, device=device) # [J] |
| 325 | suffix = W32.gather(1, idx2.expand(R, J)) # [R, J] |
| 326 | K = N - 1 # offsets 1..N-1 |
| 327 | chunk = max(1, _CHUNK_BYTE_BUDGET // (_CHUNK_PEAK_BYTES_PER_ELEM * max(1, R * J))) |
| 328 | |
| 329 | # amax keeps one penalty per token, for its longest match, however the offsets |
| 330 | # are chunked, as long as the penalty does not fall as the match grows |
| 331 | # (base >= 1, mult >= 0). |
| 332 | pen = torch.zeros(R * vocab + 1, dtype=torch.float32, device=device) |
| 333 | |
| 334 | for k0 in range(1, K + 1, chunk): |
| 335 | k1 = min(k0 + chunk, K + 1) |
| 336 | ks = torch.arange(k0, k1, device=device) # [C] |
| 337 | C = ks.shape[0] |
| 338 | # idx1[c, j] = N-1-j-k ; invalid (out of window) entries masked. |
| 339 | idx1 = idx2[None, :] - ks[:, None] # [C, J] |
| 340 | invalid = idx1 < (N - n_r)[:, None, None] # [R, C, J] |
| 341 | eq = W32[:, idx1.clamp(min=0)] == suffix[:, None, :] # [R, C, J] |
| 342 | eq &= ~invalid |
| 343 | del invalid |
| 344 | # Run length of leading True along j = the match length. Explicit dtypes avoid |
| 345 | # int64 intermediate copies; values are 0/1 and runs are <= J <= _J_BUDGET, |
| 346 | # so int8/int16 are exact. |
| 347 | L = eq.cumprod(dim=2, dtype=torch.int8).sum(dim=2, dtype=torch.int16) |
| 348 | del eq |
| 349 | # Do not narrow rep_limit to L's int16; window lengths can exceed it. |
| 350 | L = torch.minimum(L, rep_limit[:, None]) |
| 351 | |
| 352 | # Follower token of offset k lives at column N-k; the offset counts only |
| 353 | # while position i = N-1-k is inside the window (k <= n_r - 1). |
| 354 | valid_k = ks[None, :] <= (n_r - 1)[:, None] # [R, C] |
| 355 | charge = (allowed[:, None] <= L) & valid_k & active[:, None] |
| 356 | # No `if charge.any()` and no boolean indexing: both sync the host. Uncharged |
| 357 | # entries scatter to a trash slot one past the accumulator end, never read. |
| 358 | followers = W.gather(1, (N - ks).clamp(max=N - 1).expand(R, C)) |
| 359 | rows = torch.arange(R, device=device)[:, None].expand(R, C) * vocab |
| 360 | flat = rows + followers.clamp(min=0) |
| 361 | # A Python int, not a device tensor: torch.tensor(x, device=...) here would be a |
| 362 | # pageable host-to-device copy, which is itself a synchronization. |
| 363 | flat = torch.where(charge, flat, R * vocab) |
| 364 | # L takes rep_limit's dtype from the minimum above. The cast keeps the |
| 365 | # exponent int64 whatever dtype n_r, allowed and max_exp arrive in. |
| 366 | exponent = torch.minimum(L.to(torch.int64) - allowed[:, None], max_exp[:, None]) |
| 367 | # float64 pow, as llama.cpp's std::pow(float, int); a float32 pow overflows |
| 368 | # early (0.8 * 2**128 is finite in float32). |
| 369 | p_chunk = mult[:, None].double() * torch.pow( |
| 370 | base[:, None].double(), exponent.to(torch.float64) |
| 371 | ) |
| 372 | # No where() on p_chunk: uncharged entries went to the trash slot and are |
| 373 | # never read. |
| 374 | pen.scatter_reduce_( |
| 375 | 0, |
| 376 | flat.reshape(-1), |
| 377 | p_chunk.to(torch.float32).reshape(-1), |
| 378 | reduce="amax", |
| 379 | include_self=True, |
| 380 | ) |
| 381 | |
| 382 | # Of the penalty terms, only the penalty itself is held at [R, vocab] (4 B/entry); |
| 383 | # the exponent, its float64 cast, the pow result and the product stay at [R, C] |
| 384 | # inside the loop. |
| 385 | # Breakers zero the penalty rather than the logit, so a breaker keeps its value. |
| 386 | # The narrowing to logits.dtype below is a no-op under vLLM's sampler, which |
| 387 | # always passes float32 logits. |
| 388 | pen2 = pen[:-1].view(R, vocab) |
| 389 | if bmask is not None: |
| 390 | pen2.masked_fill_(bmask, 0.0) |
| 391 | # index_add_ avoids the extra [R, vocab] gather+copy of `logits[row_idx] -= pen2`. |
| 392 | pen2.neg_() |
| 393 | logits.index_add_( |
| 394 | 0, row_idx, pen2 if logits.dtype == pen2.dtype else pen2.to(logits.dtype) |
| 395 | ) |
| 396 | return logits |
| 397 | |
| 398 | |
| 399 | # --------------------------------------------------------------------------- |
| 400 | # Breaker resolution |
| 401 | # --------------------------------------------------------------------------- |
| 402 | |
| 403 | |
| 404 | class _VocabIndex(NamedTuple): |
| 405 | """A tokenizer's decoded vocabulary, plus a character index into it. |
| 406 | |
| 407 | ``char_ids`` maps every character in ``texts`` to the ids of the tokens |
| 408 | whose text contains it. |
| 409 | """ |
| 410 | |
| 411 | texts: list[str] |
| 412 | char_ids: dict[str, list[int]] |
| 413 | |
| 414 | |
| 415 | # Per-tokenizer caches, weakly keyed so tokenizers can be collected: |
| 416 | # tokenizer -> its decoded vocabulary and character index, and |
| 417 | # tokenizer -> {breaker string tuple -> resolved breaker ids}. |
| 418 | _BreakerIdsPerTokenizer = dict[tuple[str, ...], list[int]] |
| 419 | _vocab_index_cache: "weakref.WeakKeyDictionary[Any, _VocabIndex]" = ( |
| 420 | weakref.WeakKeyDictionary() |
| 421 | ) |
| 422 | _breaker_ids_cache: "weakref.WeakKeyDictionary[Any, _BreakerIdsPerTokenizer]" = ( |
| 423 | weakref.WeakKeyDictionary() |
| 424 | ) |
| 425 | |
| 426 | |
| 427 | def _vocab_index(tokenizer: Any) -> _VocabIndex: |
| 428 | """Decode every token of ``tokenizer`` once and index its characters.""" |
| 429 | index = _vocab_index_cache.get(tokenizer) |
| 430 | if index is None: |
| 431 | # Include added tokens above vocab_size, which llama.cpp's scan also covers. |
| 432 | n_ids = getattr(tokenizer, "max_token_id", tokenizer.vocab_size - 1) + 1 |
| 433 | texts = tokenizer.batch_decode([[i] for i in range(n_ids)]) |
| 434 | char_ids: dict[str, list[int]] = {} |
| 435 | for i, text in enumerate(texts): |
| 436 | for char in set(text): |
| 437 | char_ids.setdefault(char, []).append(i) |
| 438 | index = _VocabIndex(texts, char_ids) |
| 439 | _vocab_index_cache[tokenizer] = index |
| 440 | return index |
| 441 | |
| 442 | |
| 443 | def resolve_dry_breakers(tokenizer: Any, breaker_strs: tuple[str, ...]) -> list[int]: |
| 444 | """Resolve breaker strings to the ids of every token whose text contains one. |
| 445 | |
| 446 | This is llama.cpp's ``get_overlapping_token_sequences`` rule, without its |
| 447 | multi-token restart sequences. Resolution runs in ``add_request``, inside |
| 448 | ``execute_model``, so a single-character breaker (llama.cpp's defaults and |
| 449 | most custom sets) resolves by lookup in the character index instead of a |
| 450 | vocabulary scan. Results are cached per (tokenizer, breaker set). |
| 451 | """ |
| 452 | breaker_strs = tuple(s[:_MAX_BREAKER_CHAR_LEN] for s in breaker_strs if s) |
| 453 | if not breaker_strs: |
| 454 | return [] |
| 455 | per_tok = _breaker_ids_cache.setdefault(tokenizer, {}) |
| 456 | cached = per_tok.get(breaker_strs) |
| 457 | if cached is not None: |
| 458 | return list(cached) |
| 459 | |
| 460 | index = _vocab_index(tokenizer) |
| 461 | ids: set[int] = set() |
| 462 | for s in breaker_strs: |
| 463 | if len(s) == 1: |
| 464 | ids.update(index.char_ids.get(s, ())) |
| 465 | else: |
| 466 | ids.update(i for i, text in enumerate(index.texts) if s in text) |
| 467 | result = sorted(ids) |
| 468 | if len(per_tok) >= _MAX_CACHED_BREAKER_SETS: |
| 469 | # Evict the oldest entry (dict preserves insertion order). |
| 470 | per_tok.pop(next(iter(per_tok))) |
| 471 | per_tok[breaker_strs] = result |
| 472 | return list(result) |
| 473 | |
| 474 | |
| 475 | # --------------------------------------------------------------------------- |
| 476 | # The logits processor |
| 477 | # --------------------------------------------------------------------------- |
| 478 | |
| 479 | |
| 480 | def _dry_args(sampling_params: SamplingParams) -> dict[str, Any]: |
| 481 | """Read the request's DRY arguments, filling in the defaults. |
| 482 | |
| 483 | A key present with value None counts as unset, matching how |
| 484 | ``SamplingParams`` coerces its own None-valued arguments. |
| 485 | """ |
| 486 | extra = sampling_params.extra_args or {} |
| 487 | return { |
| 488 | name: default if extra.get(name) is None else extra[name] |
| 489 | for name, default in _DRY_DEFAULTS.items() |
| 490 | } |
| 491 | |
| 492 | |
| 493 | def use_dry(multiplier: float, base: float, penalty_last_n: int) -> bool: |
| 494 | """Whether DRY applies, by llama.cpp's gate in llama_sampler_dry_apply.""" |
| 495 | return bool(multiplier) and base >= 1.0 and penalty_last_n != 0 |
| 496 | |
| 497 | |
| 498 | class DryState(LogitsProcessor): |
| 499 | def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): |
| 500 | # Refuse speculative decoding here, since admission cannot see its config. |
| 501 | if vllm_config.speculative_config is not None: |
| 502 | raise ValueError(STR_SPEC_DEC_REJECTS_DRY) |
| 503 | self.req_states = req_states |
| 504 | max_num_reqs = req_states.max_num_reqs |
| 505 | self.vocab_size = req_states.vocab_size |
| 506 | self.device = req_states.device |
| 507 | |
| 508 | # float32, to round multiplier and base as llama.cpp's float members do. |
| 509 | self.multiplier = np.zeros(max_num_reqs, dtype=np.float32) |
| 510 | self.base = np.zeros(max_num_reqs, dtype=np.float32) |
| 511 | self.allowed_length = np.zeros(max_num_reqs, dtype=np.int64) |
| 512 | self.penalty_last_n = np.zeros(max_num_reqs, dtype=np.int64) |
| 513 | self.max_exponent = np.zeros(max_num_reqs, dtype=np.int64) |
| 514 | self.use_dry = np.zeros(max_num_reqs, dtype=bool) |
| 515 | |
| 516 | # req_idx -> its breaker set, and breaker set -> [vocab] bool device |
| 517 | # mask shared by every request that asked for the same breakers. |
| 518 | self.breaker_ids: dict[int, frozenset[int]] = {} |
| 519 | self._breaker_masks: dict[frozenset[int], torch.Tensor] = {} |
| 520 | |
| 521 | self._warned_unresolved = False |
| 522 | |
| 523 | # Deferred import: the tokenizer registry pulls in transformers, and |
| 524 | # the frontend imports this module only to validate params. |
| 525 | from vllm.tokenizers import cached_tokenizer_from_config |
| 526 | |
| 527 | # None under skip_tokenizer_init. |
| 528 | self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) |
| 529 | # Resolve the default set here to keep the vocabulary decode out of |
| 530 | # execute_model, where add_request runs. |
| 531 | self._default_breaker_ids = self._resolve_breakers( |
| 532 | DEFAULT_DRY_SEQUENCE_BREAKERS |
| 533 | ) |
| 534 | self._warm_up_dry_core() |
| 535 | |
| 536 | def _warm_up_dry_core(self) -> None: |
| 537 | """Load dry_core's CUDA kernels before a request needs them. |
| 538 | |
| 539 | The operands match the dtypes and ranks apply passes, and include a breaker |
| 540 | mask so that the breaker path's kernels load too. |
| 541 | """ |
| 542 | vocab = self.vocab_size |
| 543 | # The breaker takes the last id, which must not be the window's token 0. |
| 544 | if vocab < 2: |
| 545 | return |
| 546 | device = self.device |
| 547 | allowed = _DRY_DEFAULTS["dry_allowed_length"] |
| 548 | base = _DRY_DEFAULTS["dry_base"] |
| 549 | max_exp = max_exponent(base) |
| 550 | |
| 551 | def col(value: float, dtype: torch.dtype) -> torch.Tensor: |
| 552 | return torch.full((1,), value, dtype=dtype, device=device) |
| 553 | |
| 554 | breakers = torch.zeros(vocab, dtype=torch.bool, device=device) |
| 555 | breakers.narrow(0, vocab - 1, 1).fill_(True) |
| 556 | dry_core( |
| 557 | torch.zeros(1, vocab, dtype=torch.float32, device=device), |
| 558 | row_idx=col(0, torch.int64), |
| 559 | W=torch.zeros(1, _WARMUP_WINDOW, dtype=torch.int64, device=device), |
| 560 | n_r=col(_WARMUP_WINDOW, torch.int64), |
| 561 | allowed=col(allowed, torch.int64), |
| 562 | max_exp=col(max_exp, torch.int64), |
| 563 | mult=col(1.0, torch.float32), |
| 564 | base=col(base, torch.float32), |
| 565 | breaker_masks=[breakers], |
| 566 | j_budget=allowed + max_exp, |
| 567 | ) |
| 568 | |
| 569 | @classmethod |
| 570 | def validate_params(cls, sampling_params: SamplingParams) -> None: |
| 571 | """Check the ``dry_*`` keys of ``extra_args`` at request admission. |
| 572 | |
| 573 | Raises: |
| 574 | ValueError: on an unknown or out-of-range DRY argument. |
| 575 | |
| 576 | """ |
| 577 | extra = sampling_params.extra_args or {} |
| 578 | unknown = sorted( |
| 579 | key for key in extra if key.startswith("dry_") and key not in _DRY_DEFAULTS |
| 580 | ) |
| 581 | if unknown: |
| 582 | # A misspelled key would otherwise be ignored silently. |
| 583 | raise ValueError( |
| 584 | f"Unknown dry_* extra_args: {', '.join(unknown)}. " |
| 585 | f"Supported keys: {', '.join(_DRY_DEFAULTS)}." |
| 586 | ) |
| 587 | |
| 588 | args = _dry_args(sampling_params) |
| 589 | for name in ( |
| 590 | "dry_multiplier", |
| 591 | "dry_base", |
| 592 | "dry_allowed_length", |
| 593 | "dry_penalty_last_n", |
| 594 | ): |
| 595 | value = args[name] |
| 596 | # JSON true/false are bools, which pass isinstance(value, int). |
| 597 | if isinstance(value, bool) or not isinstance(value, (int, float)): |
| 598 | raise ValueError( |
| 599 | f"{name} must be a number, got {type(value).__name__}." |
| 600 | ) |
| 601 | |
| 602 | multiplier = args["dry_multiplier"] |
| 603 | base = args["dry_base"] |
| 604 | allowed_length = args["dry_allowed_length"] |
| 605 | penalty_last_n = args["dry_penalty_last_n"] |
| 606 | breakers = args["dry_sequence_breakers"] |
| 607 | |
| 608 | if not math.isfinite(multiplier) or multiplier < 0.0: |
| 609 | raise ValueError( |
| 610 | f"dry_multiplier must be non-negative and finite, got {multiplier}." |
| 611 | ) |
| 612 | if not math.isfinite(base) or base < 0.0: |
| 613 | raise ValueError(f"dry_base must be non-negative and finite, got {base}.") |
| 614 | if ( |
| 615 | not isinstance(allowed_length, int) |
| 616 | or allowed_length < 0 |
| 617 | or allowed_length > _DRY_INT_MAX |
| 618 | ): |
| 619 | raise ValueError( |
| 620 | f"dry_allowed_length must be an integer in [0, {_DRY_INT_MAX}], " |
| 621 | f"got {allowed_length}." |
| 622 | ) |
| 623 | if ( |
| 624 | not isinstance(penalty_last_n, int) |
| 625 | or penalty_last_n < -1 |
| 626 | or penalty_last_n > _DRY_INT_MAX |
| 627 | ): |
| 628 | raise ValueError( |
| 629 | "dry_penalty_last_n must be an integer: -1 (whole context), " |
| 630 | f"0 (disable), or in [1, {_DRY_INT_MAX}], got {penalty_last_n}." |
| 631 | ) |
| 632 | if not isinstance(breakers, (list, tuple)) or any( |
| 633 | not isinstance(s, str) for s in breakers |
| 634 | ): |
| 635 | raise ValueError( |
| 636 | f"dry_sequence_breakers must be a list of strings, got {breakers!r}." |
| 637 | ) |
| 638 | if len(breakers) > MAX_DRY_SEQUENCE_BREAKERS: |
| 639 | raise ValueError( |
| 640 | f"dry_sequence_breakers supports at most " |
| 641 | f"{MAX_DRY_SEQUENCE_BREAKERS} entries, got {len(breakers)}." |
| 642 | ) |
| 643 | if multiplier and 0.0 <= base < 1.0: |
| 644 | # libllama has the same gate. llama-server resets such a base to its |
| 645 | # default instead. |
| 646 | logger.warning( |
| 647 | "dry_base=%s is below 1.0, which disables DRY entirely " |
| 648 | "(llama.cpp semantics), even though dry_multiplier=%s was " |
| 649 | "set. No repetition penalty will be applied.", |
| 650 | base, |
| 651 | multiplier, |
| 652 | ) |
| 653 | |
| 654 | def _resolve_breakers(self, breakers: tuple[str, ...]) -> frozenset[int]: |
| 655 | if self._tokenizer is None or not breakers: |
| 656 | return frozenset() |
| 657 | return frozenset(resolve_dry_breakers(self._tokenizer, breakers)) |
| 658 | |
| 659 | def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: |
| 660 | args = _dry_args(sampling_params) |
| 661 | multiplier = args["dry_multiplier"] |
| 662 | base = args["dry_base"] |
| 663 | penalty_last_n = args["dry_penalty_last_n"] |
| 664 | enabled = use_dry(multiplier, base, penalty_last_n) |
| 665 | self.use_dry[req_idx] = enabled |
| 666 | self.breaker_ids.pop(req_idx, None) |
| 667 | if not enabled: |
| 668 | return False |
| 669 | self.multiplier[req_idx] = multiplier |
| 670 | self.base[req_idx] = base |
| 671 | self.allowed_length[req_idx] = args["dry_allowed_length"] |
| 672 | self.penalty_last_n[req_idx] = penalty_last_n |
| 673 | self.max_exponent[req_idx] = max_exponent(float(self.base[req_idx])) |
| 674 | |
| 675 | breakers = tuple(args["dry_sequence_breakers"]) |
| 676 | ids = ( |
| 677 | self._default_breaker_ids |
| 678 | if breakers == DEFAULT_DRY_SEQUENCE_BREAKERS |
| 679 | else self._resolve_breakers(breakers) |
| 680 | ) |
| 681 | if ids: |
| 682 | self.breaker_ids[req_idx] = ids |
| 683 | elif breakers and self._tokenizer is None and not self._warned_unresolved: |
| 684 | logger.warning( |
| 685 | "DRY sequence breakers were not resolved to token ids: this " |
| 686 | "engine has no tokenizer. Proceeding without breakers." |
| 687 | ) |
| 688 | self._warned_unresolved = True |
| 689 | return True |
| 690 | |
| 691 | def _breaker_mask(self, req_idx: int) -> torch.Tensor | None: |
| 692 | ids = self.breaker_ids.get(req_idx) |
| 693 | if not ids: |
| 694 | return None |
| 695 | mask = self._breaker_masks.get(ids) |
| 696 | if mask is None: |
| 697 | # Built on the host to avoid a sync; once per breaker set. |
| 698 | ids_np = np.fromiter(ids, dtype=np.int64, count=len(ids)) |
| 699 | ids_np = ids_np[ids_np < self.vocab_size] |
| 700 | m_np = np.zeros(self.vocab_size, dtype=bool) |
| 701 | m_np[ids_np] = True |
| 702 | mask = async_tensor_h2d(m_np, self.device) |
| 703 | if len(self._breaker_masks) >= _MAX_CACHED_BREAKER_MASKS: |
| 704 | # Evict the oldest set. A client can send a new set with every request. |
| 705 | self._breaker_masks.pop(next(iter(self._breaker_masks))) |
| 706 | self._breaker_masks[ids] = mask |
| 707 | return mask |
| 708 | |
| 709 | def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: |
| 710 | req_indices = ctx.idx_mapping_np |
| 711 | active_rows = np.flatnonzero(self.use_dry[req_indices]) |
| 712 | if active_rows.size == 0: |
| 713 | return logits |
| 714 | if logits.shape[0] != req_indices.shape[0]: |
| 715 | raise RuntimeError("DRY received draft-expanded logits") |
| 716 | |
| 717 | # host-side bound; reading positions off GPU would sync per step |
| 718 | cur_len = ctx.seq_lens_upper_bound_np[active_rows].astype(np.int64) |
| 719 | |
| 720 | reqs = req_indices[active_rows] |
| 721 | last_n = self.penalty_last_n[reqs] |
| 722 | window_len = np.where(last_n == -1, cur_len, np.minimum(cur_len, last_n)) |
| 723 | allowed = self.allowed_length[reqs] |
| 724 | keep = window_len > allowed |
| 725 | if not np.any(keep): |
| 726 | return logits |
| 727 | active_rows = active_rows[keep] |
| 728 | reqs = reqs[keep] |
| 729 | cur_len = cur_len[keep] |
| 730 | window_len = window_len[keep] |
| 731 | allowed = allowed[keep] |
| 732 | max_exp = self.max_exponent[reqs] |
| 733 | |
| 734 | # Route degenerate-clamp requests (base <= 1.000001 or oversized |
| 735 | # cap) through the sequential reference implementation. |
| 736 | fast = (max_exp > 0) & (allowed + max_exp <= _J_BUDGET) |
| 737 | all_tokens = self.req_states.all_token_ids.gpu |
| 738 | |
| 739 | if np.any(fast): |
| 740 | f_rows = active_rows[fast] |
| 741 | f_reqs = reqs[fast] |
| 742 | f_len = window_len[fast] |
| 743 | N = int(f_len.max()) |
| 744 | reqs_t = async_tensor_h2d(f_reqs, self.device) |
| 745 | cur_t = async_tensor_h2d(cur_len[fast], self.device) |
| 746 | j = torch.arange(N, device=self.device) |
| 747 | # Right-aligned gather: column j holds token (cur_len - N + j); |
| 748 | # out-of-window columns are masked inside dry_core via n_r. |
| 749 | gather_idx = (cur_t[:, None] - N + j[None, :]).clamp(min=0) |
| 750 | W = all_tokens[reqs_t[:, None], gather_idx].long() |
| 751 | dry_core( |
| 752 | logits, |
| 753 | row_idx=async_tensor_h2d(f_rows, self.device), |
| 754 | W=W, |
| 755 | n_r=async_tensor_h2d(f_len, self.device), |
| 756 | allowed=async_tensor_h2d(allowed[fast], self.device), |
| 757 | max_exp=async_tensor_h2d(max_exp[fast], self.device), |
| 758 | mult=async_tensor_h2d(self.multiplier[f_reqs], self.device), |
| 759 | base=async_tensor_h2d(self.base[f_reqs], self.device), |
| 760 | breaker_masks=[self._breaker_mask(r) for r in f_reqs], |
| 761 | j_budget=int((allowed[fast] + max_exp[fast]).max()), |
| 762 | ) |
| 763 | |
| 764 | # The sequential fallback copies each window to the host, an expected sync. |
| 765 | slow = ~fast |
| 766 | if np.any(slow): |
| 767 | with gpu_sync_allowed(): |
| 768 | rows_list = [] |
| 769 | cols_list = [] |
| 770 | vals_list = [] |
| 771 | for row, req, w_len, cur in zip( |
| 772 | active_rows[slow], reqs[slow], window_len[slow], cur_len[slow] |
| 773 | ): |
| 774 | window = ( |
| 775 | all_tokens[int(req), int(cur) - int(w_len) : int(cur)] |
| 776 | .cpu() |
| 777 | .tolist() |
| 778 | ) |
| 779 | penalties = _dry_penalties( |
| 780 | window, |
| 781 | self.breaker_ids.get(int(req), frozenset()), |
| 782 | float(self.multiplier[req]), |
| 783 | float(self.base[req]), |
| 784 | int(self.allowed_length[req]), |
| 785 | int(self.max_exponent[req]), |
| 786 | ) |
| 787 | for tok, val in penalties.items(): |
| 788 | rows_list.append(int(row)) |
| 789 | cols_list.append(tok) |
| 790 | vals_list.append(val) |
| 791 | if rows_list: |
| 792 | logits[ |
| 793 | torch.tensor(rows_list, dtype=torch.int64, device=self.device), |
| 794 | torch.tensor(cols_list, dtype=torch.int64, device=self.device), |
| 795 | ] -= torch.tensor( |
| 796 | vals_list, dtype=torch.float32, device=self.device |
| 797 | ) |
| 798 | return logits |
| 799 | |
| 800 | |
| 801 | # --------------------------------------------------------------------------- |
| 802 | # The demo |
| 803 | # --------------------------------------------------------------------------- |
| 804 | |
| 805 | |
| 806 | def main(): |
| 807 | # Imported here so that loading DryState from this module does not import LLM. |
| 808 | from vllm import LLM |
| 809 | |
| 810 | prompts = ["The future of AI is"] * 2 |
| 811 | sampling_params_list = [ |
| 812 | SamplingParams(temperature=0.0, max_tokens=64), |
| 813 | SamplingParams( |
| 814 | temperature=0.0, max_tokens=64, extra_args={"dry_multiplier": 0.8} |
| 815 | ), |
| 816 | ] |
| 817 | llm = LLM(model="facebook/opt-125m", logits_processors=[DryState]) |
| 818 | outputs = llm.generate(prompts, sampling_params_list) |
| 819 | print("\nGenerated Outputs:\n" + "-" * 60) |
| 820 | for label, output in zip(("without DRY", "with DRY"), outputs): |
| 821 | print(f"Prompt: {output.prompt!r}, {label}") |
| 822 | print(f"Output: {output.outputs[0].text!r}") |
| 823 | print("-" * 60) |
| 824 | |
| 825 | |
| 826 | if __name__ == "__main__": |
| 827 | main() |