Logit Processor additional data

[julien-piet]

Hi,

Currently, logit processors are only aware of the previous generated tokens and the logits:

vllm/vllm/model_executor/layers/sampler.py

Line 210 in 2acd76f

logits_row = logits_processor(token_ids, logits_row)

For some applications (like watermarking), additional context can be useful, such as the index of the generation in the original batch. It would be nice if these could be passed as optional parameters to logit processors. Passing the seq_id variable would be neat for instance.

Additionally, having a higher level interface (such as a LogitProcessor similar to the one in HuggingFace that handles batches instead of individual generations) could be useful (see the changes in main...julien-piet:vllm:main).

Happy to help out with this!
Julien

[DarkLight1337]

#4985 recently added the possibility for LogitsProcessor to accept the (input) prompt token IDs.

Feel free to bring up any additional arguments to LogitsProcessor you want vLLM to support in #4769.