Unified system improves AI model choice and memory use for faster responses

Unified AI Gateway: A Framework for Joint Model Routing and KV Cache Management

Distributed, Parallel, and Cluster Computing

Summary

Using large AI language models can be tricky because there are many different models that vary in speed, cost, and skills. Switching among these models and managing their memory for previous conversations takes extra time and computing work. The authors propose a system called a Unified AI Gateway that smartly decides which model to use, where to run it, and how to handle memory to speed things up while controlling costs. It also manages memory in the background to get ready for future requests. Their analysis shows this approach can make responses much faster and cheaper in many common use cases.

large language modelmodel routingkey-value cacheinferencelatencycompute placementedge computingcache managementdistributed storage

Authors

Jiaxun Lu, Xiang Zhang, Yunfeng Shao

Abstract

Large language model (LLM) inference increasingly spans models that differ in size, capability, price, and provider. This shift creates two costs for developers. One is the integration cost of choosing among and switching between many models. The other is the inference cost of rebuilding a KV cache when it is unavailable or incompatible with the selected model. We define and analyze the Unified AI Gateway as a system setting for an edge-deployed AI traffic hub. It coordinates model routing, KV cache management, and compute placement across end devices, edge resources, and cloud model services. At request time, the gateway jointly selects a target model, an execution site, and a KV cache action under task-quality, latency, cost, and resource constraints. In parallel, background cache-management actions optimize KV cache placement, replication, retrieval, and lifecycle decisions for subsequent requests. We synthesize existing evidence on KV cache reuse, compression, cross-model mapping, distributed storage, and transfer, and discuss the remaining challenges of integrating these capabilities into one system. Across eight typical workload profiles, our workload-level analytical simulation reports TTFT speedups of 1.25$\times$--13.28$\times$ and input-cost benefits of 1.20$\times$--6.16$\times$.