Customize prefix caching

Hi everyone, I have been tinkering with LLM serving lately and I was wondering if it was possible to provide a custom implementation for prefix caching when serving a model in the mode that accepts input embeddings rather than token ids list.

I know that floating point arithmetic is tricky reason why prefix caching is used for integer tokens alone, but let’s say that I’d like to try something custom that when the endpoint receives direct input embeddings, is there a way to do so or I need to fork the engine?

As per the docs, I couldn’t find a plugin mechanism to do so.

Prefix caching for direct input embeddings (i.e., when serving models with embeddings instead of token IDs) is not natively supported in vLLM as of now. The current prefix caching mechanism is designed around token IDs, and prompt embedding support is handled separately; there is no documented plugin or extension point for custom prefix caching logic with embeddings. If you want to experiment with custom caching for embeddings, you would need to fork and modify the engine yourself, as there is no built-in API or plugin system for this use case (source, source, source).

Recent discussions confirm that prefix caching for prompt embeddings is a complex topic due to the lack of a canonical hash for floating-point tensors, and while there are ongoing efforts and RFCs, official support is not yet available. You may follow or contribute to related issues and PRs for future updates (source, source).

Would you like more detail on the technical challenges or pointers on where to start modifying the engine?

Sources:

Hint: Mention @Herald in the post for followups.