| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Sequence |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | pytest.importorskip("triton") |
| 12 | if not torch.cuda.is_available(): |
| 13 | pytest.skip("CUDA required for sampler flag tests", allow_module_level=True) |
| 14 | |
| 15 | from vllm.sampling_params import SamplingParams |
| 16 | from vllm.v1.worker.gpu.sample.logits_processor import ( |
| 17 | LogitsContext, |
| 18 | LogitsProcessor, |
| 19 | ) |
| 20 | from vllm.v1.worker.gpu.sample.sampler import Sampler |
| 21 | from vllm.v1.worker.gpu.states import RequestState |
| 22 | |
| 23 | DEVICE = torch.device("cuda") |
| 24 | VOCAB_SIZE = 128 |
| 25 | |
| 26 | |
| 27 | class MockReasoningConfig: |
| 28 | reasoning_start_token_ids = [90] |
| 29 | reasoning_end_token_ids = [91] |
| 30 | natural_reasoning_end_token_ids = [91] |
| 31 | |
| 32 | |
| 33 | def _make_sampler(custom_logits_processors: Sequence[LogitsProcessor] = ()) -> Sampler: |
| 34 | req_states = RequestState( |
| 35 | max_num_reqs=4, |
| 36 | max_model_len=64, |
| 37 | max_num_batched_tokens=16, |
| 38 | num_speculative_steps=1, |
| 39 | vocab_size=VOCAB_SIZE, |
| 40 | device=DEVICE, |
| 41 | ) |
| 42 | return Sampler( |
| 43 | vllm_config=SimpleNamespace(reasoning_config=MockReasoningConfig()), |
| 44 | max_num_reqs=4, |
| 45 | vocab_size=VOCAB_SIZE, |
| 46 | device=DEVICE, |
| 47 | req_states=req_states, |
| 48 | custom_logits_processors=custom_logits_processors, |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | @pytest.mark.parametrize( |
| 53 | ("sampling_params", "expected"), |
| 54 | [ |
| 55 | pytest.param(SamplingParams(), False, id="defaults"), |
| 56 | pytest.param(SamplingParams(temperature=0.0), False, id="greedy"), |
| 57 | pytest.param( |
| 58 | SamplingParams(thinking_token_budget=3), True, id="thinking-budget" |
| 59 | ), |
| 60 | pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"), |
| 61 | pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"), |
| 62 | pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"), |
| 63 | pytest.param(SamplingParams(temperature=0.7), True, id="temperature"), |
| 64 | pytest.param(SamplingParams(min_p=0.1), True, id="min-p"), |
| 65 | pytest.param(SamplingParams(top_k=10), True, id="top-k"), |
| 66 | pytest.param(SamplingParams(top_p=0.9), True, id="top-p"), |
| 67 | pytest.param( |
| 68 | SamplingParams.for_sampler_warmup(), True, id="all-logits-processors" |
| 69 | ), |
| 70 | ], |
| 71 | ) |
| 72 | def test_logits_processing_cache_matches_request_features( |
| 73 | sampling_params: SamplingParams, expected: bool |
| 74 | ): |
| 75 | sampler = _make_sampler() |
| 76 | sampler.add_request(3, sampling_params=sampling_params) |
| 77 | |
| 78 | assert sampler.needs_logits_processing[3] == expected |
| 79 | |
| 80 | |
| 81 | def test_logits_processing_cache_is_overwritten_when_slot_is_reused(): |
| 82 | sampler = _make_sampler() |
| 83 | sampler.add_request(3, SamplingParams.for_sampler_warmup()) |
| 84 | sampler.add_request(3, SamplingParams()) |
| 85 | |
| 86 | assert not sampler.needs_logits_processing[3] |
| 87 | |
| 88 | |
| 89 | class _GateProcessor(LogitsProcessor): |
| 90 | """Reports a fixed admission decision from add_request().""" |
| 91 | |
| 92 | def __init__(self, admitted: bool): |
| 93 | self.admitted = admitted |
| 94 | |
| 95 | def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: |
| 96 | return self.admitted |
| 97 | |
| 98 | def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: |
| 99 | raise AssertionError("the pipeline never runs in this test") |
| 100 | |
| 101 | |
| 102 | @pytest.mark.parametrize("admitted", [True, False]) |
| 103 | def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool): |
| 104 | """A custom processor's add_request() return value is OR-ed into the |
| 105 | per-request needs_logits_processing flag.""" |
| 106 | sampler = _make_sampler(custom_logits_processors=[_GateProcessor(admitted)]) |
| 107 | sampler.add_request(3, SamplingParams()) |
| 108 | |
| 109 | assert sampler.needs_logits_processing[3] == admitted |
| 110 | |
| 111 | |
| 112 | def test_logits_processing_cache_only_checks_active_requests(): |
| 113 | sampler = _make_sampler() |
| 114 | sampler.add_request(0, SamplingParams(temperature=0.0)) |
| 115 | sampler.add_request(2, SamplingParams.for_sampler_warmup()) |
| 116 | |
| 117 | sampling_only = np.array([0], dtype=np.int32) |
| 118 | with_processing = np.array([0, 2], dtype=np.int32) |
| 119 | |
| 120 | assert not np.any(sampler.needs_logits_processing[sampling_only]) |
| 121 | assert np.any(sampler.needs_logits_processing[with_processing]) |