[RFC]: Logits processor extensibility

[afeldman-nm]

Motivation.

Users want logits processor extensibility, i.e. the ability to specify logits processors beyond those such as min-p which are hard-coded into the engine. See for example:

• [Bug]: V1 engine ignores logits processors and min-p sampling #12678
• https://github.com/NVIDIA/logits-processor-zoo - library of logits processor extensions

The purpose of this RFC is to establish the interface for extending the vLLM V1 engine with additional logits processors during engine instantiation.

vLLM V0 supports logits processor configuration at request level (SamplingParams attribute). For V0 running in server mode, PR #11150 makes it possible for a request to dynamically import one or more logits processor modules, assuming that the necessary modules are available/installed. The logits_processors argument (available in the completion, chat completion and transcription API endpoints) allows the custom logits processors’ constructors to be specified as a list of (1) qualified names, or (2) LogitsProcessorConstructor data structures (which include the qualified name along with constructor arguments). For security purposes (prevention of arbitrary code execution), ​​--logits-processor-pattern whitelists specific logits processor libraries via regex.

We expect V1 will add logits processor support, with logits processors instantiated at server init time. See RFC #13360 , PR #16728 . (Note that, although the logits processors are instantiated at server init time, the behavior of the logits processors - including but not limited to enabling a given logits processor - can still be controlled using SamplingParams.extra_args on a per-request basis. #16862 allows SamplingParams.extra_args to be configured via via the vllm_xargs REST API argument.) #16728 adds a logits processor base class and migrates several hard-coded logits processors (min P, min tokens, logits bias) to be sub-classes of this base class. However, #16728 does not make the list of logit processors in a given engine instance extensible beyond the builtins - thus, this RFC focused on the need to implement extensibility as a follow-on task.

Support for logits processor extensibility in v1 is desirable, both for server mode and also for direct instantiation of LLM and AsyncLLM in Python.

Proposed Change.

Interface

• For the purpose of this workstream - which is solely considering vLLM v1 engine - a "logits processor" is a subclass of the LogitsProcessor class (as defined in vllm/v1/sample/logits_processor.py in https://github.com/vllm-project/vllm/pull/16728/files ). Note that at time of writing, third-party libraries such as logits-processor-zoo are not directly compatible with this programming model.

• In server mode, vLLM V1 engine will have a new CLI argument, logits-processors, which passes in a list of logits processor constructors. It will be necessary to find a clean way to represent this on the command line; here I propose (1) logits-processors expects a string representation of a JSON-formatted list, in which each element is one of the following: (1) qualified name of a logits processor, or (2) a "constructor" JSON object with a key/value pair for the qualified name of the logits processor as well as optional positional args (args) and keyword args (kwargs) arguments; for reference see [Frontend] Add logits_processors as an extra completion argument #11150 (comment)

  • Each logit processor module specified via command line will be imported
  • Each imported logits processor will be instantiated in the persistent batch; see InputBatch implementation in [V1] LogitsProcessor programming model #16728
  • Note that there will be no need for a logits-processor-pattern CLI argument (unlike in V0) because there is no need to whitelist specific module names, since the user is specifying logits processor names explicitly via CLI
  • Example:

# CLI logits processor example
vllm serve ... --logits-processors '['logits_processor_zoo.vllm.GenLengthLogitsProcessor',{'qualname': 'vllm.v1.sample.MinPLogitsProcessor','args':[0,'argument_value'],'kwargs':{'arg_name': 'arg_value'}}]'

• LLM and AsyncLLM engine will support a logits_processors constructor argument. The argument still accepts a list of logits processors specifications, however unlike the CLI interface, each list element may be one of (1) a logits processor subclass, (2) the qualified name of a logits processor subclass, or (3) an instance of vllm.entrypoints.openai.protocol.LogitsProcessorConstructor (the underlying Python class for the aforementioned constructor JSON objects in the CLI interface):

class LogitsProcessorConstructor(BaseModel):
    qualname: str
    args: Optional[List[Any]] = None
    kwargs: Optional[Dict[str, Any]] = None

• Users instantiating the engine directly in Python are responsible for security concerns as regards logits processors (i.e. validating that they are not importing third-party logits processor libraries in an unsafe way).
• Examples:

from vllm.entrypoints.openai.protocol import LogitsProcessorConstructor
from logits_processor_zoo.vllm import CiteFromPromptLogitsProcessor
...

logitprocs_list = [CiteFromPromptLogitsProcessor,
                   'logits_processor_zoo.vllm.GenLengthLogitsProcessor',
                   LogitsProcessorConstructor(qualname='vllm.v1.sample.MinPLogitsProcessor',
                                                             args=[0,'argument_value'],
                                                             kwargs={'arg_name': 'arg_value'})]

# Sync engine example
llm = LLM(model="facebook/opt-125m",
          logits_processors=logitprocs_list)

# Async engine example
async_llm = AsyncLLM(model="facebook/opt-125m",
                     logits_processors=logitprocs_list)

• In server mode, the logits processors specified by the logits-processors CLI argument, will be passed to the logits_processors argument of the engine constructor

• For V0 back-compatibility, V0 continues supporting logits_processors request argument in REST API until V0 is removed

• The vLLM V1 engine raises an invalid exception when V0 logits processor interfaces (SamplingParams logits_processors, REST API logits_processors) are utilized (it probably makes sense to implement this check in protocol.py even though protocol.py is not v1-specific, because we want skip importing logits processor modules if the user is mistakenly using the v0 logits processor interface with the v1 engine.)

Implementation

1. Logits processors are passed in through the server CLI interface or through the LLM/AsyncLLM constructors.
2. For logits processors which were specified by qualified names, the qualified names are resolved during engine initialization. The end result is a list of logits processor classes
3. The list of logits processors is passed to the InputBatch constructor, which instantiates each logits processor

Feedback Period.

2 weeks

CC List.

@njhill @russellb @simon-mo @WoosukKwon

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[njhill]

Thanks @afeldman-nm, from a first pass this looks reasonable to me. It would be worth aligning with how other plugins are handled in general, if applicable, see: #7131, #7426.

Something else I had in mind, which I think can also be a follow-on to the core work, is to still support "old style" LPs, via a special adapter implementation of the new LP interface. The old style LP would have to provide a factory method / constructor, then in the adapter impl we just instantiate one of them per new request and keep the dict of those as its state (the factory takes the sampling params so can pull from them as needed and return None in the case that the LP doesn't apply for that request). In the "apply" we call the per-request LPs individually for the corresponding rows of the logits tensor.

Not suggesting doing this before the core stuff is done/solid, but it should mean that existing LPs still can be used easily (e.g. from the nvidia lp zoo).

[aarnphm]

Thanks @afeldman-nm for writing this, I have a few question following up

vLLM V0 supports logits processor configuration at request level (SamplingParams attribute). For V0 running in server mode, PR #11150 makes it possible for a request to dynamically import one or more logits processor modules, assuming that the necessary modules are available/installed.

iiuc from #13360 we don't want to allow request-level logit processor, so having this at the server initialization makes sense if users want to include their own logit processors (for sending requests I think it is best to parse it from extra_body or extra_args in SamplingParams)

I would also agree with Nick's suggestions for having somewhat similar interface as plugins.

The logits_processors argument (available in the completion and chat completion API endpoints) allows the custom logits processors’ constructors to be specified as a list of (1) qualified names

I feel like most custom logit processors as a plugin should include a default where users don't have to do a lot of configuration. I don't know in which use-case where customization is needed here (correct me if I'm wrong, but having a good default is imo better)

# CLI logits processor example
vllm serve ... --logits-processors '['logits_processor_zoo.vllm.GenLengthLogitsProcessor',{'qualname': 'vllm.v1.sample.MinPLogitsProcessor','args':[0,'argument_value'],'kwargs':{'arg_name': 'arg_value'}}]'

This is somewhat not ergonomic, given that you will need to know what args and kwargs are for these. I would like to see maybe something like this:

pip install logit-processors-zoo # it can has a python entrypoint `gen_length = "logit_processors_zoo.vllm.GenLengthLogitsProcessor"` which maps args `gen_length` -> `GenLengthLogitsProcessor`.

vllm serve ... # by default it will include the GenLengthLogitsProcessor
# On client, users can do something like:
# client.chat.completions.create(..., extra_body=dict(gen_length=256)) # then this will invoke GenLengthLogitsProcessor

• LLM and AsyncLLM engine will support a logits_processors constructor argument. The argument still accepts a list of logits processors specifications, however unlike the CLI interface, each list element may be one of (1) a logits processor subclass, (2) the qualified name of a logits processor subclass, or (3) an instance of vllm.entrypoints.openai.protocol.LogitsProcessorConstructor (the underlying Python class for the aforementioned constructor JSON objects in the CLI interface):

class LogitsProcessorConstructor(BaseModel):
    qualname: str
    args: Optional[List[Any]] = None
    kwargs: Optional[Dict[str, Any]] = None

This class implies that this would be import lazily? In the case of offline inference, I think it is fine to just import the class and we can have the following API

from lib import CustomLogitProcessor

logit_processors = [
  CiteFromPromptLogitsProcessor,
  'path.GenLengthLogitsProcessor',
  LazyLogitProcessor(CustomLogitProcessor, *args, **kwargs) # you can determine the first args whether to be a class or a string to import, so we can support that cases
]                  

[chaunceyjiang]

This is somewhat not ergonomic, given that you will need to know what args and kwargs are for these. I would like to see >maybe something like this:

pip install logit-processors-zoo # it can has a python entrypoint `gen_length = >"logit_processors_zoo.vllm.GenLengthLogitsProcessor"` which maps args `gen_length` -> `GenLengthLogitsProcessor`.

vllm serve ... # by default it will include the GenLengthLogitsProcessor
# On client, users can do something like:
# client.chat.completions.create(..., extra_body=dict(gen_length=256)) # then this will invoke GenLengthLogitsProcessor

I have a question about this approach, similar to this issue: #15418 (comment). I only have a single Python file—how can I register this LogitsProcessor in that case?

[aarnphm]

I have a question about this approach, similar to this issue: #15418 (comment). I only have a single Python file—how can I register this LogitsProcessor in that case?

for logit_processor packages I think it would be this

[project.entry-points."vllm.logit_processors"]
gen_length = "logit_processor_zoo.vllm.GenLengthLogitProcessor"

for standalone files it involves a bit more hacking with sys.path, but we probably don't want to do that. In this case, I think best to create a custom python packages (should be relatively straightforward these days with uv using hatch backend a la carte)

[youkaichao]

I'm not familiar with logits processor, from the plugin system point of view, right now we have two main use cases:

• register out-of-tree models
• register out-of-tree hardware backends

In both cases, only a few plugins will be present in the environment. e.g. we expect at most one hardware plugin, and only several out-of-tree model registration plugins when model developers develop their private models. Therefore, we load all plugins by default.

It seems we have lots of logits processors, and if we use plugin system, they will all be loaded, affecting the start-time if users installed logit_processor_zoo .

Maybe we can use entry-points."vllm.logit_processors" to register logits processors, and a cli argument --allowed_logits_processors ( a list of plugin names) to control the logits processors we will load. This way, we don't need to pass construction arguments from cli.

[youkaichao]

after a second thought, I think you only need --logits_processors_getter="func.qual.name", and the function returns a dict of available logits processors.

for example: --logits_processors_getter="a.b.func" , and

# in module a.b

def func():
    return {
    "gen_length": logit_processor_zoo.vllm.GenLengthLogitProcessor(),
    "name": Cls(*args, **kwargs)
}

users can customize func to add whatever logits processors they want, with any construction arguments.

[aarnphm]

# in module a.b

def func():
    return {
    "gen_length": logit_processor_zoo.vllm.GenLengthLogitProcessor(),
    "name": Cls(*args, **kwargs)
}

users can customize func to add whatever logits processors they want, with any construction arguments.

This assumes that a.b is within PYTHONPATH (assuming that there is a/b in the current directory that runs vllm/serve, which is brittle)

For plugin if we want to optimize load time, then by logit processors should be lazily loaded by default fwiw (all we want really is to avoid side effects from 3rd library)

[youkaichao]

This assumes that a.b is within PYTHONPATH (assuming that there is a/b in the current directory that runs vllm/serve, which is brittle)

they can put it in a package, not under current directory. similar to --worker-cls and --worker-extension-cls, we can accept a qualified name as long as we can resolve it.