[core] Sampling controller interface

[mmoskal]

This patch adds SamplingController object on LLMEngine. It subsumes LogitsProcessor functions in SamplingParams as suggested in #5423. Instead of calling the class LogitPostProcessor I used SamplingController since it's a bit broader than just dealing with logits (in particular it allows for fast forward tokens, see below). I'm happy to rename if needed.

The basic idea is that the engine holds an instance of SamplingController and calls methods on it to influence the sampling process. In every step the following methods are called:

• prepare(sampling_metadata: SamplingMetadata): this is meant to start computation of logit biases for sequences in the batch; the controller will likely need to store mapping from sequences to indices in logit tensor
• model forward pass is started
• transform_logits(logits: torch.Tensor) -> torch.Tensor is called on the entire logit tensor (for all sequences in the batch)
• sampling is performed as usual
• transform_sampler_output(output: SamplerOutput) -> SamplerOutput is called on the output of the sampler

In case of an empty step (where no sequences are scheduled to run), the empty_step() method is called, instead of the three methods mentioned above. This is to allow the controller to perform cleanup.

To be clear, the only way to use this right now, is to derive from SamplingController and use vllm as a library.

The transformation of sampler output is primary useful in conjunction with the newly added SequenceOutput.fast_forward_tokens field. If this is set, these tokens are to be added to the sequence instead of the sampled token.

An example, where fast forward tokens are useful is generating data adhering to a certain JSON schema. The controller first forces {"name":" to be generated, then the model generates John", the controller forces ,\n"age":, model generates 42, and so on. Another example is chain-of-thought reasoning, where after the model generated a sentence, the controller forces more instructions for the model, the model generates more text, and so on. If used, these greatly speed up generation process.

The SamplingController is passed from LLMEngine to worker/executor via ExecuteModelRequest and ModelRunnerInputBase. Only driver worker receives the controller (as it's the only one that does sampling).

I also added request_id to SequenceGroupToSample (which is referenced from SamplingMetadata) to allow the controller to identify the sequences in the batch.

There are two places where I had to make changes to allow empty fast forward tokens:

• SequenceData.update_num_computed_tokens
• BlockTable.append_token_ids The main reason to support empty tokens is for the cases when the computation of the next token mask did not complete in time. Instead of waiting for it to finish, or aborting the sequence, the controller can instead do an "empty pass" on that particular sequence and try it again in the next step, without holding up the rest of the batch. Of course, if this happens several times, the controller is free to terminate the sequence.

Status: I have this working with AICI. I think this is good to go, though it may need some tests.

CC @rkooo567 @simon-mo @GindaChen @cadedaniel @njhill

FIX #5423

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

---

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Model] for adding a new model or improving an existing model. Model name should appear in the title.
• [Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
• [Kernel] for changes affecting CUDA kernels or other compute kernels.
• [Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
• [Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• We adhere to Google Python style guide and Google C++ style guide.
• Pass all linter checks. Please use format.sh to format your code.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
• Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

• After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
• After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
• After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
• Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!