[RFC]: Custom sampling params support in REST API

[afeldman-nm]

Update: after incorporating feedback, the updated proposal is described in this comment: #17191 (comment)

Original RFC proposal (outdated):

Motivation

Addresses #16802 (“Support custom args in OpenAI (chat) completion requests”) by adding an “extra” sampling params argument to all endpoints which trigger sampling (completion, chat and transcription). This is ultimately a prerequisite for logits processor support ( RFC: #13360 PR: #16728 ), since logits processors may require custom arguments which are not utilized by vLLM core sampling logic.

Proposed Change.

Here it is proposed that when using the HTTP client, custom sampling arguments may be passed in as key/value pairs via the extra_sampling_params argument

extra_sampling_params: Optional[dict[str, Any]]

#13300 added an extra_args member to SamplingParams

extra_args: Optional[dict[str, Any]] = None

protocol.py defines a class type for each endpoint’s requests. Currently, the arrival of a completion/chat/transcription request at a particular REST API endpoint causes a call to the to_sampling_params() method associated with an instance of the appropriate request class. This method constructs a SamplingParams instance from the request attributes using the from_optional() method; the proposed change is to pass extra_sampling_params to extra_args at that point:

SamplingParams.from_optional(..., extra_args=extra_sampling_params)

In this way, the custom arguments stored in SamplingParams.extra_args will be available to logits processors downstream in the request processing pipeline.

For example,

curl http://0.0.0.0:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{"model": "facebook/opt-125m", "prompt": "Say this is a test", “ignore_eos”: true, “extra_sampling_params”: {“custom_arg": <value>}}’

results in a SamplingParams instance with extra_args = {“custom_arg": <value>}.

This RFC only applies to API endpoints which trigger sampling, summarized below (along with their associated request classes in protocols.py):

• /v1/completions (CompletionRequest)
• /v1/chat/completions (ChatCompletionRequest)
• /v1/audio/transcriptions (TranscriptionRequest)

The following API endpoints do not trigger sampling and are not part of this workstream (note that to save time in writing this RFC, I refer to the endpoints in terms of broad categories here):

• Embeddings (EmbeddingCompletionRequest, EmbeddingChatRequest)
• Rerank (RerankRequest)
• Tokenization/Detokenization (TokenizationCompletionRequest, TokenizationChatRequest, DetokenizeRequest)
• LoRA load (LoadLoRAAdapterRequest) and unload (UnloadLoRAAdapterRequest)

If you are using the OpenAI Python SDK (or similar SDK in another language), the client-side completion/chat/transcription request method does not have an extra_sampling_params argument; extra_sampling_params will need to be passed in as a key-value pair to the extra_body dict argument of the request method. Note that the extra_body argument is not part of the server’s REST API and if you pass extra_body as an argument within an HTTP client request, the server will ignore it. extra_body is simply a “catch-all” argument supported by the Python SDK to handle “special” parameters. Internally, the SDK unpacks extra_body into REST API arguments. The server does not see the extra_body argument.

Under the proposed changes in this PR, the following SDK request exemplifies a correct usage:

   completion = await client.completions.create(model=model_name,
                                                prompt="Hello, my name is",
                                                max_tokens=5,
                                                temperature=0.0,
                                                extra_body={“ignore_eos”: True,
                                                “extra_sampling_params”:           
                                                {“custom_arg”: True})

• OpenAI-standard API arguments are set directly as arguments to create()
• Arguments such as ignore_eos are set in extra_body but not in extra_sampling_params, because ignore_eos is an argument defined explicitly in protocols.py and utilized by vLLM’s core sampling functionality
• custom_arg (which is meant to represent a hypothetical custom argument for a logits processor) is not defined explicitly in any of the request types defined in protocol.py and is therefore packed within extra_sampling_params

Plan for rolling out extra sampling params:

PR #16862 is WIP and does not yet satisfy the specifications below, but will by the time it lands

• In protocol.py, add an extra_sampling_params member to CompletionRequest, ChatCompletionRequest, and TranscriptionRequest.
• In each of these three request classes, extra_sampling_params is assigned to SamplingParams.extra_args inside of the to_sampling_params() method as described above.
• This PR is a prerequisite for near-term work on logits processor support.
• This PR does not introduce breaking changes.

Thoughts on alternative proposals

The core requirement is that custom sampling arguments are supported, in order to enable the logits processor workstream.

However, in discussions about the API surface area for sampling arguments, one additional proposal was that sampling arguments such as ignore_eos which are not part of the OpenAI API specification, but which are part of the core vLLM sampling implementation (i.e. they are not “custom” logits processor arguments), should be grouped together under a catch-all dict argument (perhaps under extra_sampling_params, or perhaps under a separate dict argument). In other words these would not be top-level arguments, which is currently the case if you use the HTTP client.

Here I suggest that this would add little benefit other than strict-er compliance with the OpenAI API specification, and in exchange would add unnecessary complexity and code changes.

Feedback Period.

1 week

CC List.

@njhill @comaniac @WoosukKwon @simon-mo

CC @robertgshaw2-redhat

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.

[russellb]

An opaque passthrough for use by custom logits processors makes sense to me.

One thing we need to be careful of is (at least IMO) we should NOT use an opaque field like this for parameters that are supported by code included in vLLM. I think we are better served by declaring stronger (and clearer) types for that. Perhaps during implementation, you could leave some comments along these lines to encourage others in the future not to abuse the field.

[russellb]

Something I think we should consider here and for any future items we're adding to the body of an OpenAI request: We should namespace them somehow to avoid possible conflicts with official parameters added by OpenAI in the future. For example, instead of extra_sampling_params, do something with a vllm prefix, like vllm_sampling_params.

[njhill]

Thanks @afeldman-nm I think this is good but am still a bit unsure. Having a nested extra_sampling_params in the request json is pretty verbose/cumbersome, e.g. for simple curl examples etc. We could mitigate that by:

• Using a shorter name e.g. just "extras". TBH I think "sampling_params" in the name is kind of superfluous anyhow.
• Instead of nested, have custom args be top-level but with a rule that their names must begin with e.g. x_, so that they won't conflict with any future built-in args.

For the same reason (as well as for backwards-compatibility of course), I don't think we should move the built-in vllm-specific params. Though I also understand @russellb's point about this.

[russellb]

I'm not suggesting moving any existing params, but I think namespacing our custom params is the right thing to do. We don't own that namespace, openai does. They could step on anything we're using at any time. Namespacing would avoid that.

[russellb]

Having a nested extra_sampling_params in the request json is pretty verbose/cumbersome

but it should be rare, right? It would only be relevant to custom out-of-tree code (or it should be)

[afeldman-nm]

Based on feedback I propose the following RFC revision:

• Inside CompletionRequest, ChatCompletionRequest, and TranscriptionRequest, extra_sampling_params becomes vllm_extras for improved brevity
• Rather than making vllm_extras completely opaque (as is currently the case in my RFC proposal & as @russellb suggested), we can implement vllm_extras in a fashion that will be forward-compatible with a plugin-based logits processor interface: (1) By the time vLLM initialization is finished, the engine contain a pydantic model which embodies all valid vllm_extras arguments, and during initial request processing vllm_extras should be validated against the Pydantic model. (2) Eventually (once logits processor support is fully-realized) the Pydantic model should be assembled dynamically at vLLM startup, by aggregating required arguments from all logits processor plugins. (3) This implies that each logits processor plugin includes a dict of required vllm_extras arguments and their type signatures.
• However, do not actually implement plugin support as part of this "custom args" workstream; implement the skeleton of custom args support as part of this workstream, and then leave the plugin implementation to the "logits processor interface" workstream (RFC: [RFC][V1] LogitsProcessor interface #13360 ). More detail below.

Regarding the process for logits processor plugins to extend the pydantic model that validates vllm_extras, here is an example of how a Pydantic model can be extended with additional arguments:

PluginField = Tuple[type, Any]

class UserModel(BaseModel):
    id: int
    name: str
    email: str

plugin_fields: Dict[str, PluginField] = {
    "age": (int, Field(..., ge=0, description="User age must be ≥ 0")),
    "favourite_colour": (str, Field(..., min_length=1)),
}

create_model(
        f"{core.__name__}WithPlugins",
        __base__=UserModel,
        **plugin_fields, 
    )

I like this approach because vllm_extras arguments are validated but also extensible which gives us an easy way to prevent core vllm args from being passed in.

That being said, we are not quite yet at the point of having the logits processor plugin interface fully defined and implemented (separate PR), and I believe that all of the "core" logits processors which @njhill defined in #13360 rely on core vLLM arguments such as min_p which are not part of this "custom sampling params" workstream. So for the purposes of landing custom sampling params soon & decoupling from the logits processor interface workstream, I'll probably add a fixed CustomArgsModel pydantic model with no members to the core vLLM code and vLLM will validate vllm_extras against this model. Of course this means that in the near term vLLM API will only accept an empty vllm_extras dict. However, for the purposes of unit-testing the custom args feature, it should be possible to use monkeypatch.setattr() or something similar to inject fake custom args into CustomArgsModel.

Once the logits processor interface is fully fleshed out, then it can be part of that workstream to replace the empty CustomArgsModel definition with a dynamically-constructed model derived from plugins.

Would love to know what you think of this and if this leaves any concerns unaddressed?

CC: @njhill @russellb @simon-mo but also open to feedback from anyone else

[russellb]

some thoughts

• This is specific to the custom logits processors plugin interface, right? If so, perhaps a more desscriptive name would be appropriate. Instead of vllm_extras, vllm_logits_extras or something -- I don't have strong opinions on the exact name, just that it's clear what it applies to and resides in the vllm namespace.

• Would multiple logits processor plugins be supported at the same time? If so, how is one (or more) selected? How would you deal with the different plugins having different schemas for its inputs?

Another way to deal with the validation side would be to add a validation callback as part of the plugin interface. Let the plugins deal with it and developers/users of multiple plugins can also deal with conflicts between the validations. I don't think core vllm needs to work overly hard at that.

[afeldman-nm]

Based on feedback I propose the following RFC revision:

• Inside CompletionRequest, ChatCompletionRequest, and TranscriptionRequest, extra_sampling_params becomes vllm_extras for improved brevity
• Rather than making vllm_extras completely opaque (as is currently the case in my RFC proposal & as @russellb suggested), we can implement vllm_extras in a fashion that will be forward-compatible with a plugin-based logits processor interface: (1) By the time vLLM initialization is finished, the engine contain a pydantic model which embodies all valid vllm_extras arguments, and during initial request processing vllm_extras should be validated against the Pydantic model. (2) Eventually (once logits processor support is fully-realized) the Pydantic model should be assembled dynamically at vLLM startup, by aggregating required arguments from all logits processor plugins. (3) This implies that each logits processor plugin includes a dict of required vllm_extras arguments and their type signatures.
• However, do not actually implement plugin support as part of this "custom args" workstream; implement the skeleton of custom args support as part of this workstream, and then leave the plugin implementation to the "logits processor interface" workstream (RFC: [RFC][V1] LogitsProcessor interface #13360 ). More detail below.

Regarding the process for logits processor plugins to extend the pydantic model that validates vllm_extras, here is an example of how a Pydantic model can be extended with additional arguments:

PluginField = Tuple[type, Any]

class UserModel(BaseModel):
    id: int
    name: str
    email: str

plugin_fields: Dict[str, PluginField] = {
    "age": (int, Field(..., ge=0, description="User age must be ≥ 0")),
    "favourite_colour": (str, Field(..., min_length=1)),
}

create_model(
        f"{core.__name__}WithPlugins",
        __base__=UserModel,
        **plugin_fields, 
    )

I like this approach because vllm_extras arguments are validated but also extensible which gives us an easy way to prevent core vllm args from being passed in.

That being said, we are not quite yet at the point of having the logits processor plugin interface fully defined and implemented (separate PR), and I believe that all of the "core" logits processors which @njhill defined in #13360 rely on core vLLM arguments such as min_p which are not part of this "custom sampling params" workstream. So for the purposes of landing custom sampling params soon & decoupling from the logits processor interface workstream, I'll probably add a fixed CustomArgsModel pydantic model with no members to the core vLLM code and vLLM will validate vllm_extras against this model. Of course this means that in the near term vLLM API will only accept an empty vllm_extras dict. However, for the purposes of unit-testing the custom args feature, it should be possible to use monkeypatch.setattr() or something similar to inject fake custom args into CustomArgsModel.

Once the logits processor interface is fully fleshed out, then it can be part of that workstream to replace the empty CustomArgsModel definition with a dynamically-constructed model derived from plugins.

Would love to know what you think of this and if this leaves any concerns unaddressed?

CC: @njhill @russellb @simon-mo but also open to feedback from anyone else

Alright, closing comment on this RFC since it has been greater than a week. Thanks everyone for the feedback.

I plan to move forward with the approach described in my previous message, with a few changes. Notes on the changes:

• In the REST API interface, the custom args argument will be vllm_xargs
• In the near-term, the type signature for this custom args argument will be dict[str,CustomArgValueType] where CustomArgValueType=Any. Having an Any value type does mean that Pydantic will not be able to validate the custom arg value; however I do not think we should be concerned about this for the following reasons: (1) args passed to vllm_xargs are not solely destined for logits processors; they may also (for example) be utilized by custom samplers within vllm plugins. Therefore it is not possible to anticipate the type. (2) There is precedent in protocol.py for having dict arguments with Any for the value type; for example, chat_template_kwargs: Optional[dict[str, Any]] in EmbeddingChatRequest.
• In the longer term, within a given engine instance, it may be possible to have logits processors and vllm plugins register custom args & create a pydantic model that whitelists specific custom args with defined types. This whitelist would be constructed at engine initialization, and at request time it would allow the vLLM server frontend to validate the key/value pairs in vllm_xargs against the aforementioned whitelist. However such functionality is for a future PR and is contingent on defining the process for adding custom logits processors in vLLM v1.