[RFC]: Revise Logits Processor Programming Model

[afeldman-nm]

Motivation.

The purpose of this work is to hide details of the vLLM engine implementation from logits processor implementors; currently the vLLM logits processor programming model requires the logits processor implementor to consider complex bookkeeping details about persistent batch state changes. The proposed changes decrease the amount of bookkeeping related to batch ordering which the implementor must consider, by taking advantage of upcoming changes to the vLLM model runner.

This change impacts

1. How existing builtin logits processors based on LogitsProcessor base class (min-p, logits bias, and min tokens) are implemented
2. How new builtin logits processors would be implemented
3. How custom logits processors would be implemented

There are still some logits processors hard-coded into the vLLM engine implementation - this interface change is a prerequisite for porting these hard-coded logits processors to become subclasses of LogitsProcessor.

Current logits processor programming model

Currently, to define a logits processor, you must subclass vllm.v1.sample.logits_processor.LogitsProcessor and define (at minimum) the following methods:

• __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)
  • vllm_config: engine configuration data structure
  • device: hardware accelerator device info
  • is_pin_memory: flag indicating whether pin memory is available to support logits processor implementation
• apply(self, logits: torch.Tensor) -> torch.Tensor:
  • Consume a (num_requests) x (vocab_size) logits tensor (logits)
  • Apply logits processor transformation at batch granularity
  • Return a transformed (num_requests) x (vocab_size) logits tensor
  • You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
• is_argmax_invariant(self) -> bool:
  • Return True if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), False if the logits processor may modify argmax
  • is_argmax_invariant() is evaluated once at startup; if True, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling
• update_state(self, batch_update: Optional["BatchUpdate"]) -> None:
  • Consume a BatchUpdate data structure representing persistent batch state changes at the beginning of the current engine step
  • Use the BatchUpdate members to update logits processor internal state
  • Note: batch update data structure may be None, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated output_token_ids lists that it could have retained when they were added.

These methods are also the interface presented to the engine for invoking the logits processor. So currently, the logits processor interface to the vLLM engine is the same as the programming model for defining a logits processor.

Problems

1. Currently, when the logits processor implementor writes update_state(self, batch_update: Optional["BatchUpdate"]) -> None, they must handle complex bookkeeping resulting from changes in vLLM persistent batch state. Every add, removed or reordered request in a given step potentially results in an update to logits processor internal state. These batch state details are not conceptually related to logits processors, but must be considered because vLLM logits processors operate at batch granularity.

2. Relatedly - there is a WIP revision to the vLLM model runner ( GPU Model Runner V2 #25266 ) which will simplify representation of persistent batch state and eliminate the complex bookkeeping - however, it will require a change in the logits processor engine interface.

Proposed Change.

• GPU Model Runner V2 #25266 will eliminate batch "reordering" because the persistent batch state is unordered; each step applies a new batch order on-the-fly (call this "stepwise batch ordering"). Thus the update_state() method in the logits processor interface must accept a mapping from persistent batch index (call this "persistent index") to stepwise index
• Separate "add request" and "remove request" functionality out from update_state(), to improve separation of concerns

Supporting data structures

Type annotations for added and removed requests:

# Persistent index of removed request
RemovedRequest = int

# (params, prompt length, persistent index) tuples for new
# requests added to the batch
AddedRequest = tuple[SamplingParams, int, int]

These two type annotations will be used in new logits processor add_request(), remove_request() methods

The BatchUpdate data structure is no longer necessary.

Defining a logits processor

To define a logits processor, you would subclass vllm.v1.sample.logits_processor.LogitsProcessor and define (at minimum) the following methods:

• __init__(...), is_argmax_invariant(...): unchanged
• apply(self, logits: torch.Tensor) -> torch.Tensor:
  • Essentially unchanged from current specification. However note that the mapping from logits rows to requests is determined by the current stepwise ordering.
  • Thus update_state() must leverage its index_mapping argument such that apply() knows how to transform each row of logits
• add_request(self, req_info: AddedRequest) -> None
  • Possible outcomes: (1) add info about a new request to the logits processor. (2) no change.
• remove_request(self, req_pdx: RemoveRequest) -> None
  • Remove information about a specified request from logits processor internal state, if such information is present
• update_state(self, index_mapping: Optional[npt.NDArray], token_ids: npt.NDArray, seq_lens: npt.NDArray) -> None:
  • index_mapping: optional stepwise request ordering, as a 1D numpy array with length equal to the size of the persistent batch. None means the stepwise ordering is unchanged from the previous step
  • token_ids: a (persistent batch size) x (max num tokens) numpy array. Each row index corresponds to the persistent index of a request in the persistent batch; the contents of the row are the prompt ids and output ids concatenated for that request
  • seq_lens: a 1D numpy array with length equal to the size of the persistent batch. Each index corresponds to the persistent index of a request in the persistent batch; the value at each index is the total sequence length (prompt + output tokens)
  • Note: If index_mapping is None, the LogitsProcessor might still want to update its state based on the updated token_ids or seq_lens.

Why this new approach is beneficial:

While the logits processor implementor still has to consider details of batch ordering, it is much easier to apply what is effectively a single permutation matrix (index_mapping), rather than a complex sequence of add/remove/reorder operations as was done previously.

Feedback Period.

1 Week

CC List.

@njhill @WoosukKwon @simon-mo @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.

[njhill]

Thanks @afeldman-nm. I have been thinking about this some more, especially in the context of #25266.

Here's my current thoughts, let me know what you think:

• Instead of "owning" the dict, we can just leave that to the LP after all
• The update_states method which is called every step can have the following args:
  • batch_update: optional, similar to the current one but without moved or batch_size fields. Also AddedRequest can just be the SamplingParams, without the token id lists, but maybe with a prompt length field.
  • index_mapping: optional - None means unchanged from last step. A 1d numpy array whose length is the size of the current batch and each entry is corresponding request index (i.e. the indices used in the batch update)
  • token_ids: a 2d numpy array of token ids (prompt ids + output ids in each row). The row for a given request is its (pre-mapped) request index.
  • seq_lens: 1d numpy array similar to token_ids but just containing the total prompt + output token counts for each request. Output token count can be determined by subtracting the prompt length (obtained from AddedRequest).
• apply method unchanged

[njhill]

@afeldman-nm or perhaps instead of the first arg we could have separate add_request and remove_request methods. I think that would result in a "nicer" interface.

[afeldman-nm]

Thanks @afeldman-nm. I have been thinking about this some more, especially in the context of #25266.

Here's my current thoughts, let me know what you think:

• Instead of "owning" the dict, we can just leave that to the LP after all

• The update_states method which is called every step can have the following args:

  • batch_update: optional, similar to the current one but without moved or batch_size fields. Also AddedRequest can just be the SamplingParams, without the token id lists, but maybe with a prompt length field.
  • index_mapping: optional - None means unchanged from last step. A 1d numpy array whose length is the size of the current batch and each entry is corresponding request index (i.e. the indices used in the batch update)
  • token_ids: a 2d numpy array of token ids (prompt ids + output ids in each row). The row for a given request is its (pre-mapped) request index.
  • seq_lens: 1d numpy array similar to token_ids but just containing the total prompt + output token counts for each request. Output token count can be determined by subtracting the prompt length (obtained from AddedRequest).

• apply method unchanged

Updated to reflect feedback from @njhill

[njhill]

Thanks @afeldman-nm, I think this looks good! Some minor comments:

• What are the two ints in AddRequest? I think one of them is the prompt length, but what is the other? We could consider having that be the list of prompt token ids instead, since I think we'll have that handy anyhow. Just thinking aloud on this...
• I think with this simplification we could probably eliminate the AddRequest and RemoteRequest types altogether, just have these reflected in as args of the respective methods?

[afeldman-nm]

Hi @njhill thanks for the feedback. Three things -

1. Regarding your first question - one of the ints in AddedRequest is prompt length, the other is pre-mapping batch index (what I am calling "persistent index"); without the persistent index, how will we have a lookup index into index_mapping?

2. As for your other question - including the prompt token ids is also fine - that is currently one of the fields of AddedRequest. I guess this would mean that the update_states() token_ids field would be out_token_ids (i.e. only output token ids)? And perhaps seq_lens would become out_seq_lens i.e. output sequence lengths?

3. This leads into another consideration - I think you suggested that token_ids (or maybe it would be out_token_ids, per the previous point) should be a 2D numpy array. Presumably the first dimension of this numpy array would be batch size, but what would be the second dimension? This is a bit tricky because the sequences have different and variable lengths. One possibility is (batch size) x (length of longest sequence), in which case the shape of the array would change each step; another possibility is (batch size) x (max permitted sequence length), which would be a large and very underutilized array but has the benefit that its size doesn't change each step and so you could reuse the same array object. A third option is to eschew the 2D numpy array, and instead have token_ids be a 1D array of concatenated token vectors, where the starting index of a given token vector can be found by the cumulative sum of seq_lens (this is how variable length sequences are represented in some of the attention kernels for example). Thoughts on the best approach? I guess we can break this down into two questions: (1) 2D array versus 1D array of concatenated vectors; (2) should we pre-allocate the array to the maximum possible size it could have and then reuse the same array data structure, or should we recreate the array in each step and make it as small as possible while still having space for all of the sequences? Interested in your thoughts

[njhill]

Hi @njhill thanks for the feedback. Three things -

1. Regarding your first question - one of the ints in AddedRequest is prompt length, the other is pre-mapping batch index (what I am calling "persistent index"); without the persistent index, how will we have a lookup index into index_mapping?

Ah of course, sorry I forgot this. I guess it would make sense for that to be the first field/arg then.

2. As for your other question - including the prompt token ids is also fine - that is currently one of the fields of AddedRequest. I guess this would mean that the update_states() token_ids field would be out_token_ids (i.e. only output token ids)? And perhaps seq_lens would become out_seq_lens i.e. output sequence lengths?

I don't think it will change what's in the token_ids numpy array, i.e. that will also contain the full seq/prompt tokens. This would be just for convenience but I'm now swaying back to just keeping the count ... to keep it simpler. We don't have to make a final decision on that necessarily, let's say it will either be count or list.

3. This leads into another consideration - I think you suggested that token_ids (or maybe it would be out_token_ids, per the previous point) should be a 2D numpy array. Presumably the first dimension of this numpy array would be batch size, but what would be the second dimension? This is a bit tricky because the sequences have different and variable lengths. One possibility is (batch size) x (length of longest sequence), in which case the shape of the array would change each step; another possibility is (batch size) x (max permitted sequence length), which would be a large and very underutilized array but has the benefit that its size doesn't change each step and so you could reuse the same array object. A third option is to eschew the 2D numpy array, and instead have token_ids be a 1D array of concatenated token vectors, where the starting index of a given token vector can be found by the cumulative sum of seq_lens (this is how variable length sequences are represented in some of the attention kernels for example). Thoughts on the best approach? I guess we can break this down into two questions: (1) 2D array versus 1D array of concatenated vectors; (2) should we pre-allocate the array to the maximum possible size it could have and then reuse the same array data structure, or should we recreate the array in each step and make it as small as possible while still having space for all of the sequences? Interested in your thoughts

I was thinking it will be option 2. This is because we already maintain this (very big) array internally and will just pass that same object each time. I agree it would be too big to materialize each time (and if we were doing that then better for it to be some more dense representation like you suggest). The variable lengths are handled via the seq_lens array.

I assume you're good with my second bullet?

[afeldman-nm]

Sounds good @njhill . I'm fine with the second bullet.