| 1 | # Custom Logits Processors |
| 2 | |
| 3 | This directory contains examples demonstrating how to use custom logits processors with vLLM's offline inference API. Logits processors allow you to modify the model's output distribution before sampling, enabling controlled generation behaviors like token masking, constrained decoding, and custom sampling strategies. |
| 4 | |
| 5 | ## Scripts |
| 6 | |
| 7 | ### `custom.py` — Engine-level logits processor |
| 8 | |
| 9 | Demonstrates how to instantiate vLLM with a custom logits processor class that operates at the batch level. The example uses a `DummyLogitsProcessor` that masks out all tokens except a specified `target_token` when passed via `SamplingParams.extra_args`. |
| 10 | |
| 11 | ```bash |
| 12 | python examples/features/logits_processor/custom.py |
| 13 | ``` |
| 14 | |
| 15 | ### `custom_req.py` — Request-level logits processor wrapper |
| 16 | |
| 17 | Shows how to wrap a request-level logits processor (which operates on individual requests) to be compatible with vLLM's batch-level logits processing interface. |
| 18 | |
| 19 | ```bash |
| 20 | python examples/features/logits_processor/custom_req.py |
| 21 | ``` |
| 22 | |
| 23 | ### `custom_req_init.py` — Request-level processor with engine config |
| 24 | |
| 25 | A special case of wrapping a request-level logits processor where the processor needs access to engine configuration or model metadata during initialization (e.g., vocabulary size, tokenizer info). |
| 26 | |
| 27 | ```bash |
| 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`. |
| 42 | - **`SamplingParams.extra_args`**: Use this to pass custom keyword arguments to your logits processor on a per-request basis (e.g., `target_token`). |
| 43 | - **`DummyLogitsProcessor`**: A reference implementation available in `vllm/test_utils.py` that can be used as a starting point for custom processors. |
| 44 | |
| 45 | ## Further Reading |
| 46 | |
| 47 | - [vLLM Sampling Parameters](https://docs.vllm.ai/en/latest/api/inference_params.html#sampling-parameters) |
| 48 | - [vLLM LLM API](https://docs.vllm.ai/en/latest/api/offline_inference/llm.html) |