@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a96ee5935d

[MRV2] Add DRY as a custom logits processor example (#50584)

Signed-off-by: Oleksandr Zakharchuk <oleksandr.zakharchuk@gmail.com>

sashko-zakharchuk committed Sep 29, 2026 at 21:46 UTC a96ee5935d769ab72ca5c4dea5ee9b9964c9a45e
3 files changed +1866
examples/features/logits_processor/README.md
+8
@@ -28,6 +28,14 @@ A special case of wrapping a request-level logits processor where the processor
28 python examples/features/logits_processor/custom_req_init.py
29 ```
30
31 +### `dry.py` — DRY repetition penalty (Model Runner V2)
32 +
33 +Implements the DRY (Don't Repeat Yourself) repetition penalty, ported from llama.cpp, as `DryState`, a custom logits processor for the Model Runner V2 interface. Requests turn it on through `SamplingParams.extra_args`, for example `{"dry_multiplier": 0.8}`. The example sends one prompt twice in one batch, once without DRY and once with it, and prints both outputs. `tests/v1/sample/test_dry.py` imports the processor from this file.
34 +
35 +```bash
36 +python examples/features/logits_processor/dry.py
37 +```
38 +
39 ## Key Concepts
40
41 - **Batch-level vs. request-level**: vLLM processes logits at the batch level for efficiency. If you have a per-request processor, you need to wrap it using the patterns shown in `custom_req.py` and `custom_req_init.py`.
examples/features/logits_processor/dry.py new
+827
@@ -0,0 +1,827 @@
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()
tests/v1/sample/test_dry.py new
+1031
@@ -0,0 +1,1031 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the DRY (Don't Repeat Yourself) repetition penalty."""
4 +
5 +import math
6 +import random
7 +import subprocess
8 +import sys
9 +from pathlib import Path
10 +from types import SimpleNamespace
11 +from typing import Any
12 +
13 +import numpy as np
14 +import pytest
15 +import torch
16 +
17 +import vllm.v1.worker.gpu.sample.logits_processor.loader as loader
18 +from examples.features.logits_processor import dry
19 +from examples.features.logits_processor.dry import (
20 + _MAX_CACHED_BREAKER_MASKS,
21 + DEFAULT_DRY_SEQUENCE_BREAKERS,
22 + DryState,
23 + _dry_penalties,
24 + use_dry,
25 +)
26 +from examples.features.logits_processor.dry import max_exponent as _max_exponent
27 +from vllm.exceptions import VLLMValidationError
28 +from vllm.sampling_params import SamplingParams
29 +
30 +DEVICE = torch.device("cpu")
31 +DRY_FQCN = "examples.features.logits_processor.dry:DryState"
32 +
33 +
34 +def _params(**dry_args) -> SamplingParams:
35 + """SamplingParams carrying DRY arguments the way a client sends them."""
36 + return SamplingParams(extra_args=dry_args or None)
37 +
38 +
39 +def _penalties(window, multiplier=1.0, base=2.0, allowed_length=1, breakers=()):
40 + return _dry_penalties(
41 + window,
42 + frozenset(breakers),
43 + multiplier,
44 + base,
45 + allowed_length,
46 + _max_exponent(base),
47 + )
48 +
49 +
50 +def test_worked_example():
51 + # llama.cpp's commented example (src/llama-sampler.cpp): tokens
52 + # "a b c c b c y a b c" -> per-position repeat counts
53 + # [0,0,3,1,0,2,0,0,0,0]; extenders: c (len 3), b (len 1), y (len 2).
54 + # With multiplier=1, base=2, allowed_length=1: c -> 4, b -> 1, y -> 2.
55 + got = _penalties([1, 2, 3, 3, 2, 3, 9, 1, 2, 3])
56 + assert got == {3: 4.0, 2: 1.0, 9: 2.0}
57 +
58 +
59 +def test_z_box_equality_boundary():
60 + # The Z-algorithm's copy-vs-extend equality boundary: when
61 + # cnt[last - p] == right_part_len the match may extend past the Z-box
62 + # and must be re-scanned, not copied. Copying here drops
63 + # token 0's penalty on this input.
64 + got = _penalties([0, 0, 0, 1, 0, 0], allowed_length=2)
65 + assert got == {0: 1.0, 1: 1.0}
66 +
67 +
68 +def test_exponent_clamp_float32_semantics():
69 + # llama.cpp computes FLOAT_MAX_LOG / log(base) in float32. At base=2.0
70 + # that is exactly 128.0; float64 gives 127.99999998 -> 127. The clamp
71 + # must land on 128 or long-repeat penalties come out half of
72 + # llama.cpp's and never saturate to -inf.
73 + assert _max_exponent(2.0) == 128
74 + assert _max_exponent(1.75) == 158
75 + assert _max_exponent(1.0) == 0 # disabled guard
76 + assert _max_exponent(1.0000005) == 0 # llama.cpp's 1.000001 guard
77 +
78 +
79 +def test_long_repeat_saturates_to_neg_inf():
80 + # 200 identical tokens at base=2.0: clamped exponent 128 ->
81 + # 2**128 overflows float32 -> logit must land at exactly -inf,
82 + # matching llama.cpp. Only the tensor path saturates, on the cast into
83 + # dry_core's float32 accumulator.
84 + state, all_tokens = _make_v2_state(max_num_reqs=1, vocab=16)
85 + state.add_request(
86 + 0,
87 + _params(
88 + dry_multiplier=1.0,
89 + dry_base=2.0,
90 + dry_allowed_length=0,
91 + dry_sequence_breakers=[],
92 + ),
93 + )
94 + all_tokens[0, :200] = 7
95 + logits = torch.zeros(1, 16, device=DEVICE)
96 + _apply(state, logits, np.array([0]), np.array([200]))
97 + assert logits[0, 7].item() == -float("inf")
98 +
99 +
100 +def test_allowed_length_threshold():
101 + # Match of length 2 is penalized at allowed_length=2, not at 3.
102 + window = [1, 2, 3, 1, 2]
103 + assert _penalties(window, allowed_length=3) == {}
104 + assert _penalties(window, allowed_length=2) == {3: 1.0}
105 +
106 +
107 +def test_breakers_cap_and_exclude():
108 + # An unbroken run is penalized; making that token a breaker
109 + # collapses rep_limit to 0 and produces nothing.
110 + assert _penalties([7] * 6, allowed_length=2) == {7: 8.0}
111 + assert _penalties([7] * 6, allowed_length=2, breakers=[7]) == {}
112 +
113 +
114 +def test_short_windows_are_safe():
115 + assert _penalties([]) == {}
116 + assert _penalties([5]) == {}
117 +
118 +
119 +class _StubTokenizer:
120 + """Minimal tokenizer for breaker-resolution tests."""
121 +
122 + _TEXTS = ["a", "\n", "b:\n", "hello", ":", "x:y", "*", "**"]
123 +
124 + vocab_size = len(_TEXTS)
125 +
126 + def batch_decode(self, ids_lists):
127 + return [self._TEXTS[ids[0]] for ids in ids_lists]
128 +
129 +
130 +def test_breaker_containment_resolution():
131 + # Resolution must match every token whose text contains the breaker
132 + # string (llama.cpp get_overlapping_token_sequences), not only exact
133 + # encodings.
134 + from examples.features.logits_processor.dry import resolve_dry_breakers
135 +
136 + tok = _StubTokenizer()
137 + assert resolve_dry_breakers(tok, ("\n",)) == [1, 2]
138 + assert resolve_dry_breakers(tok, (":",)) == [2, 4, 5]
139 + assert resolve_dry_breakers(tok, ("*",)) == [6, 7]
140 + assert resolve_dry_breakers(tok, ()) == []
141 + # Repeat resolution is consistent, and callers get their own copy
142 + # (mutating a result must not poison the cache).
143 + first = resolve_dry_breakers(tok, ("\n",))
144 + first.append(999)
145 + assert resolve_dry_breakers(tok, ("\n",)) == [1, 2]
146 +
147 +
148 +def test_validate_params_rejects_json_booleans():
149 + for name in (
150 + "dry_multiplier",
151 + "dry_base",
152 + "dry_allowed_length",
153 + "dry_penalty_last_n",
154 + ):
155 + with pytest.raises(ValueError, match=name):
156 + DryState.validate_params(_params(**{name: True}))
157 +
158 +
159 +def test_validate_params():
160 + for bad in (
161 + {"dry_multiplier": -1.0},
162 + {"dry_multiplier": float("nan")},
163 + {"dry_penalty_last_n": -7},
164 + {"dry_penalty_last_n": 2.5},
165 + {"dry_allowed_length": -1},
166 + {"dry_allowed_length": 1.5},
167 + {"dry_sequence_breakers": [1, 2]},
168 + {"dry_sequence_breakers": [f"b{i}" for i in range(65)]},
169 + ):
170 + with pytest.raises(ValueError):
171 + DryState.validate_params(_params(**bad))
172 + # Valid corner values pass, as does a request that says nothing at all.
173 + DryState.validate_params(SamplingParams())
174 + DryState.validate_params(_params())
175 + DryState.validate_params(_params(dry_multiplier=0.0))
176 + DryState.validate_params(_params(dry_multiplier=1.0, dry_penalty_last_n=-1))
177 + DryState.validate_params(_params(dry_multiplier=1.0, dry_sequence_breakers=[]))
178 + # An explicit null means "unset", as it does for every SamplingParams
179 + # field; it must not be read as a non-numeric value and rejected.
180 + DryState.validate_params(_params(dry_multiplier=None, dry_base=None))
181 +
182 +
183 +def test_validate_params_rejects_unknown_dry_keys():
184 + """An unknown dry_* key fails the request instead of being ignored."""
185 + with pytest.raises(ValueError, match="dry_multipliers"):
186 + DryState.validate_params(_params(dry_multiplier=0.8, dry_multipliers=0.8))
187 + # Keys outside the namespace belong to somebody else; leave them alone.
188 + DryState.validate_params(_params(dry_multiplier=0.8, target_token=3))
189 +
190 +
191 +def test_validate_params_raises_plain_valueerror():
192 + """validate_params raises ``ValueError``, as the V2 interface specifies."""
193 + with pytest.raises(ValueError) as excinfo:
194 + DryState.validate_params(_params(dry_multiplier=-1.0))
195 + assert type(excinfo.value) is ValueError
196 +
197 +
198 +def _oracle(window, breakers, multiplier, base, allowed_length, max_exponent):
199 + """llama_sampler_dry_apply's steps, with a brute-force match scan in place of
200 + the Z-algorithm that ``_dry_penalties`` ports."""
201 + m = len(window)
202 + rep_limit = m
203 + for i in range(m):
204 + if window[m - 1 - i] in breakers:
205 + rep_limit = i
206 + break
207 + if rep_limit < allowed_length:
208 + return {}
209 + cnt = [0] * m
210 + for i in range(m - 1):
211 + length = 0
212 + while (
213 + length < m
214 + and i - length >= 0
215 + and window[i - length] == window[m - 1 - length]
216 + ):
217 + length += 1
218 + cnt[i] = min(length, rep_limit)
219 + max_tok: dict[int, int] = {}
220 + for i in range(m - 1):
221 + if cnt[i] >= allowed_length:
222 + tok = window[i + 1]
223 + if max_tok.get(tok, -1) < cnt[i]:
224 + max_tok[tok] = cnt[i]
225 + penalties = {}
226 + for tok, repeat_len in max_tok.items():
227 + if tok in breakers:
228 + continue
229 + exponent = repeat_len - allowed_length
230 + if max_exponent and exponent > max_exponent:
231 + exponent = max_exponent
232 + penalties[tok] = multiplier * (base**exponent)
233 + return penalties
234 +
235 +
236 +def test_differential_against_oracle():
237 + rng = random.Random(0)
238 + for case in range(400):
239 + if case % 10 == 0:
240 + # Long, highly repetitive: the only way to reach the clamp.
241 + alphabet, n = rng.randint(1, 2), rng.randint(200, 320)
242 + else:
243 + alphabet, n = rng.randint(2, 6), rng.randint(0, 48)
244 + window = [rng.randrange(alphabet) for _ in range(n)]
245 + multiplier = rng.choice([0.5, 0.8, 1.0, 2.0])
246 + base = rng.choice([1.1, 1.75, 2.0, 3.0])
247 + allowed = rng.randint(0, 4)
248 + breakers = frozenset(
249 + rng.sample(range(alphabet), rng.randint(0, min(2, alphabet)))
250 + )
251 + max_exp = _max_exponent(base)
252 + want = _oracle(window, breakers, multiplier, base, allowed, max_exp)
253 + got = _dry_penalties(window, breakers, multiplier, base, allowed, max_exp)
254 + assert set(want) == set(got), f"case {case}: keys {want} vs {got}"
255 + for tok in want:
256 + a, b = want[tok], got[tok]
257 + assert math.isclose(a, b, rel_tol=1e-9), (
258 + f"case {case} token {tok}: {a} vs {b}"
259 + )
260 +
261 +
262 +def test_double_pow_saturation():
263 + # llama.cpp's std::pow(float, int) computes in double: 0.8 * 2**128
264 + # is finite in float32 (-2.72e38) while 1.0 * 2**128 saturates to
265 + # -inf. A float32 pow gets the first case wrong.
266 + for mult, expect_inf in ((0.8, False), (1.0, True)):
267 + state, all_tokens = _make_v2_state(max_num_reqs=1, vocab=16)
268 + state.add_request(
269 + 0,
270 + _params(
271 + dry_multiplier=mult,
272 + dry_base=2.0,
273 + dry_allowed_length=0,
274 + dry_sequence_breakers=[],
275 + ),
276 + )
277 + all_tokens[0, :200] = 7
278 + logits = torch.zeros(1, 16, device=DEVICE)
279 + _apply(state, logits, np.array([0]), np.array([200]))
280 + val = logits[0, 7].item()
281 + if expect_inf:
282 + assert val == -float("inf")
283 + else:
284 + assert math.isfinite(val) and math.isclose(
285 + val, -0.8 * 2.0**128, rel_tol=1e-6
286 + )
287 +
288 +
289 +def test_dry_core_rejects_j_budget_over_the_budget():
290 + """dry_core raises on a j_budget over _J_BUDGET and accepts _J_BUDGET itself."""
291 +
292 + def run(j_budget):
293 + def col(value, dtype=torch.int64):
294 + return torch.full((1,), value, dtype=dtype, device=DEVICE)
295 +
296 + dry.dry_core(
297 + torch.zeros(1, 8, device=DEVICE),
298 + row_idx=col(0),
299 + W=torch.zeros(1, 4, dtype=torch.int64, device=DEVICE),
300 + n_r=col(4),
301 + allowed=col(2),
302 + max_exp=col(1),
303 + mult=col(1.0, torch.float32),
304 + base=col(2.0, torch.float32),
305 + breaker_masks=[None],
306 + j_budget=j_budget,
307 + )
308 +
309 + run(dry._J_BUDGET)
310 + with pytest.raises(ValueError, match="j_budget"):
311 + run(dry._J_BUDGET + 1)
312 +
313 +
314 +# ---------------------------------------------------------------------------
315 +# The DryState custom logits processor
316 +# ---------------------------------------------------------------------------
317 +
318 +
319 +def _make_vllm_config(speculative_config=None, skip_tokenizer_init=True):
320 + """Stub VllmConfig with only the fields DryState reads."""
321 + return SimpleNamespace(
322 + speculative_config=speculative_config,
323 + model_config=SimpleNamespace(skip_tokenizer_init=skip_tokenizer_init),
324 + )
325 +
326 +
327 +def _make_v2_state(
328 + max_num_reqs=8,
329 + vocab=32,
330 + max_model_len=256,
331 + device=DEVICE,
332 + tokenizer=None,
333 +):
334 + all_tokens = torch.zeros(
335 + max_num_reqs, max_model_len, dtype=torch.int32, device=device
336 + )
337 + req_states = SimpleNamespace(
338 + max_num_reqs=max_num_reqs,
339 + vocab_size=vocab,
340 + device=device,
341 + all_token_ids=SimpleNamespace(gpu=all_tokens),
342 + )
343 + state = DryState(_make_vllm_config(), req_states)
344 + if tokenizer is not None:
345 + # Re-resolve the default set now that a tokenizer is installed.
346 + state._tokenizer = tokenizer
347 + state._default_breaker_ids = state._resolve_breakers(
348 + DEFAULT_DRY_SEQUENCE_BREAKERS
349 + )
350 + return state, all_tokens
351 +
352 +
353 +def _apply(state, logits, idx_mapping_np, seq_lens_np):
354 + """Call ``DryState.apply`` with a LogitsContext built as the sampler builds it.
355 +
356 + Args:
357 + state: the ``DryState`` under test.
358 + logits: [num_reqs, vocab] tensor, one row per request, modified in place.
359 + idx_mapping_np: [num_reqs] batch position -> request slot.
360 + seq_lens_np: [num_reqs] host-side context lengths.
361 +
362 + """
363 + from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
364 +
365 + n_reqs = idx_mapping_np.shape[0]
366 + assert logits.shape[0] == n_reqs
367 + idx_mapping = torch.from_numpy(idx_mapping_np).to(logits.device)
368 + ctx = LogitsContext(
369 + expanded_idx_mapping=idx_mapping,
370 + idx_mapping=idx_mapping,
371 + idx_mapping_np=idx_mapping_np,
372 + expanded_local_pos=torch.zeros(n_reqs, dtype=torch.int64, device=logits.device),
373 + input_ids=torch.zeros(n_reqs, dtype=torch.int32, device=logits.device),
374 + pos=torch.from_numpy(seq_lens_np.astype("int64") - 1).to(logits.device),
375 + seq_lens_upper_bound_np=seq_lens_np,
376 + )
377 + return state.apply(logits, ctx)
378 +
379 +
380 +def test_v2_worked_example_and_window_bounds():
381 + # The window is [0, seq_len), including the last input token. Without it every
382 + # penalty lands on the continuation of the previous suffix instead of the one
383 + # being sampled.
384 + state, all_tokens = _make_v2_state()
385 + params = _params(
386 + dry_multiplier=1.0,
387 + dry_base=2.0,
388 + dry_allowed_length=1,
389 + dry_sequence_breakers=[],
390 + )
391 + state.add_request(3, params)
392 + seq = [1, 2, 3, 3, 2, 3, 9, 1, 2, 3]
393 + all_tokens[3, : len(seq)] = torch.tensor(seq, dtype=torch.int32)
394 +
395 + logits = torch.zeros(1, 32, device=DEVICE)
396 + idx_mapping = np.array([3])
397 + # last input token sits at position len(seq)-1
398 + seq_lens = np.array([len(seq)])
399 + _apply(state, logits, idx_mapping, seq_lens)
400 + assert logits[0, 3].item() == -4.0
401 + assert logits[0, 2].item() == -1.0
402 + assert logits[0, 9].item() == -2.0
403 +
404 +
405 +def test_v2_disabled_and_gating():
406 + state, _ = _make_v2_state()
407 + assert not state.add_request(0, SamplingParams()) # dry off by default
408 + assert not state.add_request(1, _params(dry_multiplier=1.0, dry_base=0.5))
409 + assert not state.add_request(2, _params(dry_multiplier=1.0, dry_penalty_last_n=0))
410 + assert not state.use_dry[:3].any()
411 + # Re-adding a slot with DRY off after one with DRY on must clear it.
412 + on = _params(dry_multiplier=0.8, dry_sequence_breakers=[])
413 + assert state.add_request(4, on)
414 + assert state.use_dry[4]
415 + assert not state.add_request(4, SamplingParams())
416 + assert not state.use_dry[4]
417 +
418 +
419 +def test_v2_penalty_last_n_window():
420 + # A cap smaller than the history must limit the visible window.
421 + state, all_tokens = _make_v2_state()
422 + params = _params(
423 + dry_multiplier=1.0,
424 + dry_base=2.0,
425 + dry_allowed_length=1,
426 + dry_penalty_last_n=4,
427 + dry_sequence_breakers=[],
428 + )
429 + state.add_request(0, params)
430 + seq = [7] * 20
431 + all_tokens[0, : len(seq)] = torch.tensor(seq, dtype=torch.int32)
432 + logits = torch.zeros(1, 32, device=DEVICE)
433 + _apply(state, logits, np.array([0]), np.array([20]))
434 + # window = last 4 tokens of an identical run: longest match ending
435 + # before the suffix is 3 (rep-limited by the window), penalty 2**(3-1).
436 + assert logits[0, 7].item() == -4.0
437 +
438 +
439 +def test_spec_decode_is_refused_at_construction():
440 + """A speculative config makes DRY refuse to load."""
441 + req_states = SimpleNamespace(
442 + max_num_reqs=4,
443 + vocab_size=32,
444 + device=DEVICE,
445 + all_token_ids=SimpleNamespace(gpu=torch.zeros(4, 8, dtype=torch.int32)),
446 + )
447 + config = _make_vllm_config(speculative_config=SimpleNamespace())
448 + with pytest.raises(ValueError, match="speculative decoding"):
449 + DryState(config, req_states)
450 +
451 +
452 +def test_v2_draft_expanded_logits_raise():
453 + """``DryState.apply`` raises on more logits rows than requests.
454 +
455 + Only speculative decoding lays logits out that way, and construction refuses it.
456 + """
457 + from vllm.v1.worker.gpu.sample.logits_processor import LogitsContext
458 +
459 + state, _ = _make_v2_state()
460 + state.add_request(0, _params(dry_multiplier=0.8, dry_sequence_breakers=[]))
461 + # One request with one draft token has two logits rows.
462 + rows = torch.zeros(2, dtype=torch.int64, device=DEVICE)
463 + ctx = LogitsContext(
464 + expanded_idx_mapping=rows,
465 + idx_mapping=rows[:1],
466 + idx_mapping_np=np.array([0]),
467 + expanded_local_pos=torch.arange(2, device=DEVICE),
468 + input_ids=torch.zeros(2, dtype=torch.int32, device=DEVICE),
469 + pos=torch.tensor([7, 8], device=DEVICE),
470 + seq_lens_upper_bound_np=np.array([9]),
471 + )
472 + logits = torch.zeros(2, 32, device=DEVICE)
473 + with pytest.raises(RuntimeError, match="draft-expanded"):
474 + state.apply(logits, ctx)
475 + assert not logits.any()
476 +
477 +
478 +def test_constructs_without_speculative_config():
479 + state, _ = _make_v2_state()
480 + assert isinstance(state, DryState)
481 + assert state._tokenizer is None # skip_tokenizer_init in the stub config
482 +
483 +
484 +@pytest.mark.parametrize("chunk_budget", [4096, dry._CHUNK_BYTE_BUDGET])
485 +def test_v2_matches_reference_fuzz(chunk_budget, monkeypatch):
486 + # The V2 window-gather and routing path must agree with the sequential
487 + # reference on randomized histories, including degenerate bases that
488 + # route through the slow path. At these sizes the default budget is always one
489 + # chunk. The 4096 B case splits most calls into several, to cover how per-chunk
490 + # results combine.
491 + monkeypatch.setattr(dry, "_CHUNK_BYTE_BUDGET", chunk_budget)
492 + state, all_tokens = _make_v2_state(max_num_reqs=8, vocab=16, max_model_len=200)
493 + rng = random.Random(7)
494 + for trial in range(60):
495 + n_reqs = rng.randint(1, 5)
496 + histories = {}
497 + seq_lens = []
498 + idx_list = []
499 + for r in range(n_reqs):
500 + args: dict[str, Any] = dict(
501 + dry_multiplier=rng.choice([0.8, 1.0, 2.0]),
502 + dry_base=rng.choice([1.0000005, 1.1, 1.75, 2.0]),
503 + dry_allowed_length=rng.randint(0, 3),
504 + dry_penalty_last_n=rng.choice([-1, 5, 50]),
505 + dry_sequence_breakers=[],
506 + )
507 + state.add_request(r, _params(**args))
508 + # Breaker ids go straight into the state and must follow add_request,
509 + # which pops the slot's ids. String resolution is tested above.
510 + breakers = frozenset(rng.sample(range(4), rng.randint(0, 2)))
511 + if breakers:
512 + state.breaker_ids[r] = breakers
513 + n = rng.randint(2, 180)
514 + hist = [rng.randrange(rng.randint(1, 4) + 1) for _ in range(n)]
515 + all_tokens[r, :n] = torch.tensor(hist, dtype=torch.int32)
516 + histories[r] = (args, hist, breakers)
517 + idx_list.append(r)
518 + seq_lens.append(n)
519 +
520 + logits = torch.zeros(n_reqs, 16, device=DEVICE)
521 + _apply(state, logits, np.array(idx_list), np.array(seq_lens))
522 + for r in range(n_reqs):
523 + args, hist, breakers = histories[r]
524 + base32 = float(np.float32(args["dry_base"]))
525 + if base32 < 1.0 or not args["dry_multiplier"]:
526 + want: dict[int, float] = {}
527 + else:
528 + last_n = args["dry_penalty_last_n"]
529 + w = len(hist) if last_n == -1 else min(len(hist), last_n)
530 + allowed = args["dry_allowed_length"]
531 + window = hist[-w:] if w > allowed else None
532 + want = (
533 + _dry_penalties(
534 + window,
535 + breakers,
536 + float(np.float32(args["dry_multiplier"])),
537 + base32,
538 + allowed,
539 + _max_exponent(base32),
540 + )
541 + if window is not None
542 + else {}
543 + )
544 + for tok in range(16):
545 + expected = -want.get(tok, 0.0)
546 + got = logits[r, tok].item()
547 + assert math.isclose(expected, got, rel_tol=1e-5, abs_tol=1e-6), (
548 + f"trial {trial} req {r} tok {tok}: {expected} vs {got}"
549 + )
550 +
551 +
552 +def test_v2_batch_sharded_sampling_matches_unsharded(monkeypatch):
553 + """Each rank's DRY rows under batch-sharded sampling match the unsharded rows.
554 +
555 + With ``--enable-batch-sharded-sampling``, ``BatchSharder`` gives every TP rank
556 + a local batch of the requests it owns and localizes the two fields DRY reads
557 + from ``LogitsContext``: ``idx_mapping_np``, and ``seq_lens_cpu_upper_bound``,
558 + which the sampler passes as ``seq_lens_upper_bound_np``. Its Triton plan kernel
559 + writes only device fields DRY does not read, so it is replaced here and the
560 + sharding runs on the host.
561 + """
562 + from dataclasses import replace
563 + from unittest import mock
564 +
565 + from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
566 + from vllm.v1.worker.gpu.sample import batch_shard
567 +
568 + vocab = 16
569 + state, all_tokens = _make_v2_state(
570 + max_num_reqs=8, vocab=vocab, tokenizer=_StubTokenizer()
571 + )
572 + # Slots out of batch order, so a rank's rows are not its slots, and only rank 2 at
573 + # tp_size 3 owns a prefix of the batch. Slot 3 has DRY off, slot 6 takes the
574 + # sequential path and slot 0 caps its window. Slot 5 breaks on "\n" (tokens 1
575 + # and 2), slot 6 on "a" (token 0), and slots 7 and 1 on "b" (token 2).
576 + slots = np.array([5, 2, 7, 0, 3, 6, 1])
577 + seq_lens = np.array([40, 23, 57, 31, 12, 48, 36])
578 + breakers = {5: ["\n"], 6: ["a"], 7: ["b"], 1: ["b"]}
579 + rng = random.Random(5)
580 + for slot, n in zip(slots.tolist(), seq_lens.tolist()):
581 + args: dict[str, Any] = dict(
582 + dry_multiplier=0.8, dry_sequence_breakers=breakers.get(slot, [])
583 + )
584 + if slot == 6:
585 + args["dry_base"] = 1.0000005
586 + if slot == 0:
587 + args["dry_penalty_last_n"] = 16
588 + state.add_request(slot, SamplingParams() if slot == 3 else _params(**args))
589 + hist = [rng.randrange(3) for _ in range(n)]
590 + all_tokens[slot, :n] = torch.tensor(hist, dtype=torch.int32)
591 + assert state.breaker_ids == {5: {1, 2}, 6: {0}, 7: {2}, 1: {2}}
592 + full = torch.zeros(len(slots), vocab, device=DEVICE)
593 + _apply(state, full, slots, seq_lens)
594 + assert (full < 0).any(dim=1).tolist() == [s != 3 for s in slots.tolist()]
595 +
596 + buffers = InputBuffers(max_num_reqs=8, max_num_tokens=len(slots), device=DEVICE)
597 + batch = replace(
598 + InputBatch.make_dummy(len(slots), len(slots), buffers),
599 + req_ids=[f"req-{i}" for i in range(len(slots))],
600 + idx_mapping=torch.from_numpy(slots).to(DEVICE),
601 + idx_mapping_np=slots,
602 + seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens),
603 + )
604 + monkeypatch.setattr(batch_shard, "_build_shard_plan_kernel", mock.MagicMock())
605 + for tp_size in (2, 3):
606 + covered: list[int] = []
607 + for rank in range(tp_size):
608 + group = SimpleNamespace(rank_in_group=rank, world_size=tp_size)
609 + with mock.patch.object(batch_shard, "get_tp_group", return_value=group):
610 + sharder = batch_shard.BatchSharder(
611 + max_num_reqs=8, max_num_logits_per_req=1, device=DEVICE
612 + )
613 + local, _, _, _ = sharder.shard_sampler_inputs(batch, None)
614 + owned = [int(req_id.split("-")[1]) for req_id in local.req_ids]
615 + covered += owned
616 + logits = torch.zeros(len(owned), vocab, device=DEVICE)
617 + seq_lens_np = local.seq_lens_cpu_upper_bound.numpy()
618 + _apply(state, logits, local.idx_mapping_np, seq_lens_np)
619 + assert torch.equal(logits, full[owned]), f"tp_size {tp_size} rank {rank}"
620 + assert sorted(covered) == list(range(len(slots)))
621 +
622 +
623 +def test_v2_breaker_strings_resolve_and_reach_dry_core():
624 + # Breaker strings arrive in extra_args and are resolved by the processor's
625 + # own tokenizer; the ids must reach the penalty computation, and the slot
626 + # must not inherit them when it is reused.
627 + state, all_tokens = _make_v2_state(tokenizer=_StubTokenizer())
628 + breaking = _params(
629 + dry_multiplier=1.0,
630 + dry_base=2.0,
631 + dry_allowed_length=1,
632 + dry_sequence_breakers=["*"],
633 + )
634 + state.add_request(0, breaking)
635 + # "*" is token 6 and is contained in token 7 ("**"), so both break.
636 + assert state.breaker_ids[0] == frozenset({6, 7})
637 + all_tokens[0, :8] = 7
638 + logits = torch.zeros(1, 32, device=DEVICE)
639 + _apply(state, logits, np.array([0]), np.array([8]))
640 + assert logits[0, 7].item() == 0.0
641 +
642 + fresh = _params(
643 + dry_multiplier=1.0,
644 + dry_base=2.0,
645 + dry_allowed_length=1,
646 + dry_sequence_breakers=[],
647 + )
648 + state.add_request(0, fresh)
649 + assert 0 not in state.breaker_ids
650 + logits2 = torch.zeros(1, 32, device=DEVICE)
651 + _apply(state, logits2, np.array([0]), np.array([8]))
652 + assert logits2[0, 7].item() < 0.0
653 +
654 +
655 +def test_breaker_mask_cache_is_bounded():
656 + """The device-mask cache must not grow with the number of distinct sets."""
657 + state, _ = _make_v2_state(vocab=128)
658 + sets = [frozenset({i}) for i in range(_MAX_CACHED_BREAKER_MASKS + 8)]
659 + for breakers in sets:
660 + state.breaker_ids[0] = breakers
661 + assert state._breaker_mask(0) is not None
662 + assert len(state._breaker_masks) <= _MAX_CACHED_BREAKER_MASKS
663 + # The first set was evicted and the newest kept.
664 + assert sets[-1] in state._breaker_masks
665 + assert sets[0] not in state._breaker_masks
666 +
667 +
668 +def test_default_breaker_set_is_resolved_eagerly():
669 + """A request that names no breakers must not pay to resolve them."""
670 + state, _ = _make_v2_state(tokenizer=_StubTokenizer())
671 + # "\n" -> {1, 2}, ":" -> {2, 4, 5}, '"' -> {}, "*" -> {6, 7}.
672 + assert state._default_breaker_ids == frozenset({1, 2, 4, 5, 6, 7})
673 + state.add_request(0, _params(dry_multiplier=0.8))
674 + assert state.breaker_ids[0] is state._default_breaker_ids
675 + # Naming the default set explicitly must land on the same object.
676 + state.add_request(
677 + 1,
678 + _params(
679 + dry_multiplier=0.8,
680 + dry_sequence_breakers=list(DEFAULT_DRY_SEQUENCE_BREAKERS),
681 + ),
682 + )
683 + assert state.breaker_ids[1] is state._default_breaker_ids
684 +
685 +
686 +def test_init_warms_up_dry_core_once(monkeypatch):
687 + """__init__ calls dry_core once, the way apply calls it."""
688 + calls = []
689 + # Taken before the patch below replaces it in the same module.
690 + original = dry.dry_core
691 +
692 + def spy(logits, **kwargs):
693 + calls.append((logits, kwargs))
694 + return original(logits, **kwargs)
695 +
696 + def kernel_args(logits, kwargs):
697 + def sig(t):
698 + return None if t is None else (t.dtype, t.dim(), t.device)
699 +
700 + masks = {sig(m) for m in kwargs["breaker_masks"]}
701 + other = {
702 + name: sig(arg) if isinstance(arg, torch.Tensor) else type(arg)
703 + for name, arg in kwargs.items()
704 + if name != "breaker_masks"
705 + }
706 + return sig(logits), masks, other
707 +
708 + monkeypatch.setattr("examples.features.logits_processor.dry.dry_core", spy)
709 + _make_v2_state(vocab=1)
710 + assert not calls, "a vocabulary of one leaves no id for the breaker"
711 + state, all_tokens = _make_v2_state(vocab=64, tokenizer=_StubTokenizer())
712 + assert len(calls) == 1
713 + assert (calls[0][0] < 0).any(), "the warmup matched nothing"
714 + assert not state._breaker_masks
715 + assert not state.breaker_ids
716 + assert not state.use_dry.any()
717 +
718 + state.add_request(0, _params(dry_multiplier=0.8))
719 + all_tokens[0, :8] = 3
720 + _apply(state, torch.zeros(1, 64, device=DEVICE), np.array([0]), np.array([8]))
721 + assert len(calls) == 2
722 + assert kernel_args(*calls[0]) == kernel_args(*calls[1])
723 +
724 +
725 +def test_no_tokenizer_warns_once_and_still_penalizes():
726 + # Without a tokenizer, breaker strings are dropped with one warning.
727 + state, all_tokens = _make_v2_state()
728 + assert state._tokenizer is None
729 + assert state._default_breaker_ids == frozenset()
730 + state.add_request(0, _params(dry_multiplier=0.8))
731 + assert state._warned_unresolved
732 + assert 0 not in state.breaker_ids
733 + all_tokens[0, :8] = 7
734 + logits = torch.zeros(1, 32, device=DEVICE)
735 + _apply(state, logits, np.array([0]), np.array([8]))
736 + assert logits[0, 7].item() < 0.0
737 +
738 +
739 +def test_v2_degenerate_base_routes_to_unclamped_slow_path():
740 + # base <= 1.000001 gives max_exponent == 0, which in the sequential
741 + # reference means "no clamp"; the vectorized path would clamp the
742 + # exponent to 0, so such requests must route to the reference scan.
743 + state, all_tokens = _make_v2_state()
744 + params = _params(
745 + dry_multiplier=1.0,
746 + dry_base=1.0000005,
747 + dry_allowed_length=0,
748 + dry_sequence_breakers=[],
749 + )
750 + state.add_request(0, params)
751 + all_tokens[0, :150] = 7
752 + logits = torch.zeros(1, 32, device=DEVICE)
753 + _apply(state, logits, np.array([0]), np.array([150]))
754 + base32 = float(np.float32(1.0000005))
755 + want = -1.0 * base32**149
756 + got = logits[0, 7].item()
757 + assert math.isclose(got, want, rel_tol=1e-6)
758 + assert got != -1.0
759 +
760 +
761 +def test_breaker_resolution_covers_added_tokens():
762 + class AddedTokenTokenizer:
763 + _TEXTS = ["a", "\n", "b", "<|im_start|>"]
764 + vocab_size = 3 # the added token sits above vocab_size
765 + max_token_id = 3
766 +
767 + def batch_decode(self, ids_lists):
768 + return [self._TEXTS[ids[0]] for ids in ids_lists]
769 +
770 + from examples.features.logits_processor.dry import resolve_dry_breakers
771 +
772 + tok = AddedTokenTokenizer()
773 + assert resolve_dry_breakers(tok, ("<|im_start|>",)) == [3]
774 + assert resolve_dry_breakers(tok, ("\n",)) == [1]
775 +
776 +
777 +class _TextsTokenizer:
778 + """Tokenizer stub over an explicit list of decoded token texts."""
779 +
780 + def __init__(self, texts: list[str], vocab_size: int):
781 + self._texts = texts
782 + self.vocab_size = vocab_size
783 + self.max_token_id = len(texts) - 1
784 +
785 + def batch_decode(self, ids_lists):
786 + return [self._texts[ids[0]] for ids in ids_lists]
787 +
788 +
789 +def _containment_scan(texts: list[str], breaker_strs: tuple[str, ...]) -> list[int]:
790 + """Resolution as a plain scan: the semantics the character index must keep."""
791 + ids: set[int] = set()
792 + for s in breaker_strs:
793 + ids.update(i for i, text in enumerate(texts) if s in text)
794 + return sorted(ids)
795 +
796 +
797 +def test_breaker_index_matches_containment_scan():
798 + """Breaker resolution must match a plain containment scan.
799 +
800 + Every character of a synthetic vocabulary is checked, including multi-byte
801 + and combining ones, added tokens and an empty token, then multi-character
802 + and mixed breaker sets.
803 + """
804 + from examples.features.logits_processor.dry import (
805 + _MAX_BREAKER_CHAR_LEN,
806 + _vocab_index,
807 + resolve_dry_breakers,
808 + )
809 +
810 + rng = random.Random(11)
811 + # U+4E2D and U+1F600 are multi-byte in UTF-8. "e\u0301" is two code points, so
812 + # U+0301 alone resolves by lookup and the pair by scan.
813 + alphabet = list('abcde \n:"*.,') + ["\u4e2d", "\u00e9", "e\u0301", "\U0001f600"]
814 + texts = [""] # a token that decodes to nothing
815 + for _ in range(8000):
816 + texts.append("".join(rng.choices(alphabet, k=rng.randint(1, 12))))
817 + vocab_size = len(texts)
818 + # Added tokens, above vocab_size; one holds the vocabulary's only U+2028.
819 + texts += ["<|im_start|>", "<|im_end|>", "\u2028sep"]
820 + tok = _TextsTokenizer(texts, vocab_size)
821 +
822 + def check(breakers: tuple[str, ...]) -> list[int]:
823 + kept = tuple(s[:_MAX_BREAKER_CHAR_LEN] for s in breakers if s)
824 + want = _containment_scan(texts, kept)
825 + assert resolve_dry_breakers(tok, breakers) == want, breakers
826 + return want
827 +
828 + # Every character present is indexed, so a lookup miss means "in no token".
829 + present = sorted(set("".join(texts)))
830 + assert set(_vocab_index(tok).char_ids) == set(present)
831 + for char in present:
832 + check((char,))
833 + assert len(check(("a",))) > 2000 # a character in thousands of tokens
834 + assert check(("\u2028",)) == [len(texts) - 1] # only above vocab_size
835 + for absent in ("q", "Z", "\u2603", "\U0001f601", "\u0300"):
836 + assert check((absent,)) == []
837 +
838 + substrings = []
839 + for _ in range(300):
840 + text = rng.choice(texts[1:vocab_size])
841 + start = rng.randrange(len(text))
842 + substrings.append(text[start : start + rng.randint(2, 5)])
843 + for s in substrings + ["<|im_start|>", "e\u0301", "\u2028sep", "zz", "a" * 41]:
844 + check((s,))
845 + for _ in range(200):
846 + check(tuple(rng.sample(present, rng.randint(1, 4)) + rng.sample(substrings, 2)))
847 + assert check(()) == []
848 + assert check(("",)) == []
849 + assert check(("", "\n")) == check(("\n",))
850 +
851 +
852 +# The accumulator scales with R and the chunk transients do not, hence two batch
853 +# sizes. The breaker path's masks get their own case.
854 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
855 +@pytest.mark.parametrize(
856 + ("n_reqs", "limit_mib", "n_breakers"),
857 + [(32, 128, 0), (256, 256, 0), (256, 288, 3900)],
858 +)
859 +def test_peak_memory_bounded(n_reqs, limit_mib, n_breakers):
860 + # base=1.1 gives max_exponent=930, a large J for the vectorized path
861 + # (the cap is allowed_length + max_exponent <= _J_BUDGET = 2048).
862 + device = torch.device("cuda")
863 + rng = random.Random(3)
864 + window = 2048
865 + state, all_tokens = _make_v2_state(
866 + max_num_reqs=n_reqs, vocab=128256, max_model_len=window, device=device
867 + )
868 + params = _params(
869 + dry_multiplier=0.8,
870 + dry_base=1.1,
871 + dry_allowed_length=2,
872 + dry_sequence_breakers=[],
873 + )
874 + breaker_ids: frozenset[int] = frozenset()
875 + if n_breakers:
876 + # About the size of llama.cpp's default set on a Llama-3 tokenizer. Id 0 is
877 + # in the history, so breakers cut some matches but not all.
878 + breaker_ids = frozenset({0} | set(rng.sample(range(8, 128256), n_breakers - 1)))
879 + # Eight symbols, so 2-token repeats are common.
880 + hist = torch.tensor(
881 + [[rng.randrange(8) for _ in range(window)] for _ in range(n_reqs)],
882 + dtype=torch.int32,
883 + )
884 + all_tokens.copy_(hist.to(device))
885 + for r in range(n_reqs):
886 + state.add_request(r, params)
887 + if breaker_ids:
888 + state.breaker_ids[r] = breaker_ids
889 + logits = torch.zeros(n_reqs, 128256, device=device)
890 + torch.accelerator.synchronize()
891 + torch.accelerator.reset_peak_memory_stats()
892 + base_alloc = torch.accelerator.memory_allocated()
893 + _apply(state, logits, np.arange(n_reqs), np.full(n_reqs, window))
894 + torch.accelerator.synchronize()
895 + peak = torch.accelerator.max_memory_allocated() - base_alloc
896 + charged = int((logits < 0).sum().item())
897 + assert charged > 0, "no token was penalized"
898 + # The limits catch the exponent and pow being held at full [R, vocab] width.
899 + limit = limit_mib * 1024 * 1024
900 + assert peak < limit, f"peak {peak / 2**20:.1f} MiB over {limit / 2**20:.0f} MiB"
901 +
902 +
903 +# ---------------------------------------------------------------------------
904 +# Plugin wiring: loading by FQCN, and validation at request admission.
905 +# ---------------------------------------------------------------------------
906 +
907 +
908 +def test_loads_through_the_loader_by_fqcn():
909 + """The example's FQCN must resolve and construct."""
910 + req_states = SimpleNamespace(
911 + device=DEVICE,
912 + max_num_reqs=4,
913 + vocab_size=32,
914 + all_token_ids=SimpleNamespace(gpu=torch.zeros(4, 8, dtype=torch.int32)),
915 + prompt_len=None,
916 + prefill_len=None,
917 + total_len=None,
918 + )
919 + procs = loader.build_custom_logits_processors(
920 + _make_vllm_config(), req_states, False, [DRY_FQCN]
921 + )
922 + assert [type(p) for p in procs] == [DryState]
923 +
924 +
925 +def test_loader_import_does_not_initialize_cuda():
926 + """Importing dry.py must not initialize CUDA.
927 +
928 + The frontend imports it under ``guard_cuda_initialization()`` to validate
929 + params, and a device touched at import time fails the engine at startup.
930 + """
931 + code = (
932 + "import torch\n"
933 + "from vllm.utils.torch_utils import guard_cuda_initialization\n"
934 + "with guard_cuda_initialization():\n"
935 + " import examples.features.logits_processor.dry # noqa: F401\n"
936 + "assert not torch.cuda.is_initialized()\n"
937 + )
938 + # From the repository root, where `examples` is importable.
939 + subprocess.run(
940 + [sys.executable, "-c", code], check=True, cwd=Path(__file__).parents[3]
941 + )
942 +
943 +
944 +def test_admission_validator_rejects_bad_extra_args():
945 + """The real frontend validator must turn a bad dry_* value into a 400.
946 +
947 + ``build_custom_logits_processors_params_validator`` is what
948 + ``InputProcessor`` calls per request; it wraps a processor's ``ValueError``
949 + as ``VLLMValidationError``, which the entrypoints map to a client error.
950 + """
951 + loader._cached_load_v2_logitsprocs.cache_clear()
952 + validate = loader.build_custom_logits_processors_params_validator([DRY_FQCN])
953 + validate(SamplingParams())
954 + validate(_params(dry_multiplier=0.8))
955 + with pytest.raises(VLLMValidationError, match="dry_multiplier"):
956 + validate(_params(dry_multiplier=-1.0))
957 + with pytest.raises(VLLMValidationError, match="dry_multipliers"):
958 + validate(_params(dry_multiplier=0.8, dry_multipliers=0.8))
959 +
960 +
961 +@pytest.mark.parametrize("field", ["dry_penalty_last_n", "dry_allowed_length"])
962 +def test_dry_int_params_upper_bound(field):
963 + """Both integer fields are capped at INT32_MAX; see ``_DRY_INT_MAX``."""
964 + for bad in (2**63, 2**31):
965 + with pytest.raises(ValueError, match=field):
966 + DryState.validate_params(_params(dry_multiplier=0.8, **{field: bad}))
967 + # The boundary itself is legal, and so are the ordinary values.
968 + DryState.validate_params(_params(dry_multiplier=0.8, **{field: 2**31 - 1}))
969 + DryState.validate_params(_params(dry_multiplier=0.8, **{field: 2}))
970 +
971 +
972 +def test_dry_base_below_one_warns_that_it_disabled_dry():
973 + """`dry_base < 1.0` disables DRY, as in llama.cpp; validate_params warns.
974 +
975 + Patches the module logger because vLLM's logger does not propagate to caplog.
976 + """
977 + from unittest.mock import patch
978 +
979 + with patch("examples.features.logits_processor.dry.logger") as log:
980 + DryState.validate_params(_params(dry_multiplier=0.8, dry_base=0.8))
981 + assert not use_dry(0.8, 0.8, -1), "base < 1.0 must disable DRY"
982 + assert log.warning.called, "expected a warning"
983 + assert "disables DRY" in log.warning.call_args[0][0]
984 +
985 + # A normal configuration says nothing.
986 + with patch("examples.features.logits_processor.dry.logger") as log:
987 + DryState.validate_params(_params(dry_multiplier=0.8, dry_base=1.75))
988 + assert use_dry(0.8, 1.75, -1)
989 + assert not log.warning.called
990 +
991 +
992 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
993 +def test_result_does_not_depend_on_the_chunk_budget():
994 + """Logits must not depend on how the offsets are chunked."""
995 + from examples.features.logits_processor import dry as dc
996 +
997 + device = torch.device("cuda")
998 + vocab, window, n_reqs = 64, 600, 1
999 + # A short alphabet repeated, so one follower token recurs at offsets far
1000 + # enough apart to land in different chunks.
1001 + W = torch.tensor([([5, 6, 7] * 200)[:window]], dtype=torch.int64, device=device)
1002 + col = lambda v, dt=torch.int64: torch.full((n_reqs,), v, dtype=dt, device=device) # noqa: E731
1003 +
1004 + def run(budget):
1005 + saved = dc._CHUNK_BYTE_BUDGET
1006 + dc._CHUNK_BYTE_BUDGET = budget
1007 + try:
1008 + logits = torch.zeros(n_reqs, vocab, device=device)
1009 + dc.dry_core(
1010 + logits,
1011 + torch.arange(n_reqs, device=device),
1012 + W,
1013 + col(window),
1014 + col(2),
1015 + col(60),
1016 + col(0.8, torch.float32),
1017 + col(1.75, torch.float32),
1018 + [None] * n_reqs,
1019 + j_budget=62,
1020 + )
1021 + return logits
1022 + finally:
1023 + dc._CHUNK_BYTE_BUDGET = saved
1024 +
1025 + many = run(4096)
1026 + one = run(1 << 30)
1027 + assert torch.equal(many, one), (
1028 + f"chunking changed the result: {many[0][many[0] < 0][:3].tolist()} vs "
1029 + f"{one[0][one[0] < 0][:3].tolist()}"
1030 + )
1031 + assert (many < 0).any(), "nothing was penalized"