Research / Stanford University

Optimizing Semantic Caching for Sequence Prediction and Memory-Optimized LLM Inference Scheduling

Aakash Mishra

Stanford University · Jun 2024

PDFCodeCite
Prompts flow through an encoder and semantic cache to a token-length predictor and a scheduler in front of the vLLM engine
Prompts flow through an encoder and semantic cache to a token-length predictor and a scheduler in front of the vLLM engine

Predicts how long each LLM answer will be by reusing predictions for similar past prompts, then serves short jobs first in vLLM. 60% faster end-to-end.

Abstract

The growing prevalence of LLMs has presented a need for scalable inference platforms. Proxy-modelling for output token length presents an opportunity to schedule inference jobs more effectively using shortest-job first. We show that this method can be further optimized by using a semantic cache that uses query similarity to cache the output token length of similar prompts.

We show that the accuracy of the proxy model can be boosted with semantic caching as well as a reduction in the regressor overhead. We integrate our approach with vLLM's memory-optimized inference platform and show a 60% improvement in end-to-end vLLM inference time for the facebook/opt-125m model.

Cite

@techreport{mishra2024semantic,
  title = {Optimizing Semantic Caching for Sequence Prediction and Memory-Optimized LLM Inference Scheduling},
  author = {Aakash Mishra},
  year = {2024},
  institution = {Stanford University}
}