| 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(), |
| 921 | req_states, |
| 922 | False, |
| 923 | [DRY_FQCN], |
| 924 | ) |
| 925 | assert [type(p) for p in procs] == [DryState] |
| 926 | |
| 927 | |
| 928 | def test_loader_import_does_not_initialize_cuda(): |
| 929 | """Importing dry.py must not initialize CUDA. |
| 930 | |
| 931 | The frontend imports it under ``guard_cuda_initialization()`` to validate |
| 932 | params, and a device touched at import time fails the engine at startup. |
| 933 | """ |
| 934 | code = ( |
| 935 | "import torch\n" |
| 936 | "from vllm.utils.torch_utils import guard_cuda_initialization\n" |
| 937 | "with guard_cuda_initialization():\n" |
| 938 | " import examples.features.logits_processor.dry # noqa: F401\n" |
| 939 | "assert not torch.cuda.is_initialized()\n" |
| 940 | ) |
| 941 | # From the repository root, where `examples` is importable. |
| 942 | subprocess.run( |
| 943 | [sys.executable, "-c", code], check=True, cwd=Path(__file__).parents[3] |
| 944 | ) |
| 945 | |
| 946 | |
| 947 | def test_admission_validator_rejects_bad_extra_args(): |
| 948 | """The real frontend validator must turn a bad dry_* value into a 400. |
| 949 | |
| 950 | ``build_custom_logits_processors_params_validator`` is what |
| 951 | ``InputProcessor`` calls per request; it wraps a processor's ``ValueError`` |
| 952 | as ``VLLMValidationError``, which the entrypoints map to a client error. |
| 953 | """ |
| 954 | loader._cached_load_v2_logitsprocs.cache_clear() |
| 955 | validate = loader.build_custom_logits_processors_params_validator([DRY_FQCN]) |
| 956 | validate(SamplingParams()) |
| 957 | validate(_params(dry_multiplier=0.8)) |
| 958 | with pytest.raises(VLLMValidationError, match="dry_multiplier"): |
| 959 | validate(_params(dry_multiplier=-1.0)) |
| 960 | with pytest.raises(VLLMValidationError, match="dry_multipliers"): |
| 961 | validate(_params(dry_multiplier=0.8, dry_multipliers=0.8)) |
| 962 | |
| 963 | |
| 964 | @pytest.mark.parametrize("field", ["dry_penalty_last_n", "dry_allowed_length"]) |
| 965 | def test_dry_int_params_upper_bound(field): |
| 966 | """Both integer fields are capped at INT32_MAX; see ``_DRY_INT_MAX``.""" |
| 967 | for bad in (2**63, 2**31): |
| 968 | with pytest.raises(ValueError, match=field): |
| 969 | DryState.validate_params(_params(dry_multiplier=0.8, **{field: bad})) |
| 970 | # The boundary itself is legal, and so are the ordinary values. |
| 971 | DryState.validate_params(_params(dry_multiplier=0.8, **{field: 2**31 - 1})) |
| 972 | DryState.validate_params(_params(dry_multiplier=0.8, **{field: 2})) |
| 973 | |
| 974 | |
| 975 | def test_dry_base_below_one_warns_that_it_disabled_dry(): |
| 976 | """`dry_base < 1.0` disables DRY, as in llama.cpp; validate_params warns. |
| 977 | |
| 978 | Patches the module logger because vLLM's logger does not propagate to caplog. |
| 979 | """ |
| 980 | from unittest.mock import patch |
| 981 | |
| 982 | with patch("examples.features.logits_processor.dry.logger") as log: |
| 983 | DryState.validate_params(_params(dry_multiplier=0.8, dry_base=0.8)) |
| 984 | assert not use_dry(0.8, 0.8, -1), "base < 1.0 must disable DRY" |
| 985 | assert log.warning.called, "expected a warning" |
| 986 | assert "disables DRY" in log.warning.call_args[0][0] |
| 987 | |
| 988 | # A normal configuration says nothing. |
| 989 | with patch("examples.features.logits_processor.dry.logger") as log: |
| 990 | DryState.validate_params(_params(dry_multiplier=0.8, dry_base=1.75)) |
| 991 | assert use_dry(0.8, 1.75, -1) |
| 992 | assert not log.warning.called |
| 993 | |
| 994 | |
| 995 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 996 | def test_result_does_not_depend_on_the_chunk_budget(): |
| 997 | """Logits must not depend on how the offsets are chunked.""" |
| 998 | from examples.features.logits_processor import dry as dc |
| 999 | |
| 1000 | device = torch.device("cuda") |
| 1001 | vocab, window, n_reqs = 64, 600, 1 |
| 1002 | # A short alphabet repeated, so one follower token recurs at offsets far |
| 1003 | # enough apart to land in different chunks. |
| 1004 | W = torch.tensor([([5, 6, 7] * 200)[:window]], dtype=torch.int64, device=device) |
| 1005 | col = lambda v, dt=torch.int64: torch.full((n_reqs,), v, dtype=dt, device=device) # noqa: E731 |
| 1006 | |
| 1007 | def run(budget): |
| 1008 | saved = dc._CHUNK_BYTE_BUDGET |
| 1009 | dc._CHUNK_BYTE_BUDGET = budget |
| 1010 | try: |
| 1011 | logits = torch.zeros(n_reqs, vocab, device=device) |
| 1012 | dc.dry_core( |
| 1013 | logits, |
| 1014 | torch.arange(n_reqs, device=device), |
| 1015 | W, |
| 1016 | col(window), |
| 1017 | col(2), |
| 1018 | col(60), |
| 1019 | col(0.8, torch.float32), |
| 1020 | col(1.75, torch.float32), |
| 1021 | [None] * n_reqs, |
| 1022 | j_budget=62, |
| 1023 | ) |
| 1024 | return logits |
| 1025 | finally: |
| 1026 | dc._CHUNK_BYTE_BUDGET = saved |
| 1027 | |
| 1028 | many = run(4096) |
| 1029 | one = run(1 << 30) |
| 1030 | assert torch.equal(many, one), ( |
| 1031 | f"chunking changed the result: {many[0][many[0] < 0][:3].tolist()} vs " |
| 1032 | f"{one[0][one[0] < 0][:3].tolist()}" |
| 1033 | ) |
| 1034 | assert (many < 0).any(), "nothing was penalized" |