Lightning Talk: Not All Tokens Are Equal: Semantic KV-Cache for Agen... Maroon Ayoub & Hyunkyun Moon

This talk explores the next set of optimizations in KV cache-centric inference, building on concepts discussed in the previous session. The speakers highlight the significance of a distributed inference platform designed for heterogeneous GPU data centers, focusing on cross-vendor PDD aggregation between Nvidia and AMD technologies. They explain the importance of automatic prefix caching for KV cache reuse to maximize hardware utilization during inference workloads. The speaker introduces the challenges of state blindness in current inference systems and proposes a new layer of KV cache orchestration that involves workload awareness. By defining retention and eviction policies through a declarative language, the speakers illustrate how to manage KV cache across a cluster effectively. Examples from academic papers demonstrate the need for priority in cache management for optimal performance during multi-agent workflows, ultimately aiming to enhance KV cache orchestration.

Okay. Hi all again. Uh Marun IBM research an LMD core maintainer >> and Hung an ML pl engineer from Mor. >> Yes. So in the last talk we talked we discussed Kvash centric inference and how you're pretty much building everything around KVach and KVach orchestration. This talk takes it a step further or like what's uh the next set of optimizations we see around kvash centric uh inference.

Uh before we dive in uh I a key word for about the model that our company we build distributed inference platform for heterogeneous GPU data center and our focus our our core focus is cross vendor PDD. aggregation like in prep on Nvidia and decoder and AMD. We also do the SOD driven automatic chip selection and deployment and vertical optimization for AMD GPU clusters. Okay, so just a

slight rerun through uh for everyone who missed the previous talk on KVACcentric inference and the state of inference today which is request centric or request aware uh inference. So you have automatic prefix caching which is a key concept uh for KV cache reuse. you compute uh your context once and try and try to reuse it as much as possible so that your GPUs or your accelerators are

busy computing new state and avoiding recomputation and duplication of work as much as possible and it it works around that like the entire stack works around maximizing the effectiveness of this across your uh accelerator. So in PDD code disagregation we separate KV cache production from consumption uh scaling and specializing each uh differently. in wide expert parallelism. So it goes way beyond that but you're practically also reducing

the model weight footprint allowing for uh bigger KV cache per GPU and minimal KV cache uh duplication especially with optimizations such as data parallelism when you're working with attention such as MLA for deepseek or uh grouped query attention and so on. So effectively you're uh uh increasing the total KV cache size and availability per GPU. In intelligent inference scheduling you're uh unifying the disagregated KV cache pool

and treating it as one so that you're reusing across the cluster and not uh missing opportunities there. And in KV cache management, you're practically extending the KV cache pool. So you have uh well beyond GPU so that you have both uh bigger opportunities for KV cache storing and reuse uh therefore and with peer-to-peer sharing and storage sharing you're getting you're trying to achieve compute once reuse anywhere.

And these are all around kevcash uh state management and they they do extremely well uh for uh request centric inference where you're not really looking into the requests you're not really looking into the bigger picture or like the the agentic program being run for example and that's the current gap in in inference uh systems as we see it. So the next step really is work with our

inference. How do you better orchestrate KV cache to both maximize reuse and maximize hardware utilization by becoming uh workload aware and we'll be focusing on the agentic workload uh being uh you know the main AI workload uh pretty much at this point. So the problem really is uh state blindness. So you're not aware of what's actually inside the workload or like what's coming next or what's happening

beyond just having tokens filling up your cache. So you treat all tokens equally. So if you have a system prompt that is shared in a in an agentic workflow or in a workload. You have tool definitions that are going to be repeated over and over. And you have the conversation history with with things like uh the agentic thinking scratch pad or a very ephemeral set of tokens

that the agent we know the agent is not going to reuse. But to your system they're all the same. You're going to be uh applying LRU for eviction. So you're saving them all and you're evicting the least recently used. And that makes sense maybe for one workload, but once you have once you scale out and you have multiple agents running, you have many concurrent users, LRU becomes

suboptimal for management. So if you're able to uh retain the system prompt or retain the sections you know are going to be reused and evict even newer but uh portions you know are ephemeral you can get uh much better utilization of your uh distributed KV cache and we see that in we'll see that in some uh research papers but the main goal of this work is to

really define a new layer of KV cache orchestration where you in a declarative language you'll be able to uh define how you want KV cache to be orchestrated and following LMD's guides and wid paths concept you'd have a set of u uh out ofthe-box policies that take your inference to the next step and take your KV cache reuse uh to the next step we can cover some

we'll cover some of the policies but in very very high level So you'll be controlling you'll be using semantic context types language you understand uh to set KV cache retention policies offloading policies and movement policies across the cluster. some examples. >> Yes, here's the first example. The paper called the continum is focuses on the agent behavioral pattern. An agent regions cause a two with a two to

three seconds 30 seconds for the leard then legions again. The problem is that during weight B blocks look too old to L value. So it evicts them. Even though system prompt is uh and two definitions are guaranteed to need on the very next turn. So the core idea is uh differential retention context that guarantees to be reused every turn gets high priority like that CC and transient

context get low priority but with detail to survive the expected two post like those one. So and next this is second policy. Uh there are so paper called KB flow addresses is multiple agent workflow where agents are connected in sequence or per example is an investment analysis paralyze a pipeline. A researcher agent investigate and a financial analyst anal agent analyze and an invest recommendation agent produce a

final output. The problem is that the all three agent KB blocks compete for the same cache on the ALU. The researchers block were generated first. So they get evicted first. Even though the analyst needs the researchers output as an input text input context. The core idea is computing priority from the workflow deck. The imminent agent gets the highest priority and priority obtains dynamic dynamically as the workflows

progress. When the researchers finish the analyst becomes the imminent agent. Uh this is third policy. Uh first this pattern has three characterist characteristics. First the share the preix is very large 30k to 100k tokens. It extreme extremely valuable because all branches share it. Uh second branch specific context is transient. It has value while active but drops of once evaluated. Third branches learn in parallel but result needs

to be gathered moving to the next stop. Uh what make it this policy is special is that uh it exercise both dimension retention and priority and tail control retention tier of 14 controls movement. The previous two policies were about the retention but this one uh demonstrate for the full service our framework provide. >> Okay. So we're about to run out of time. So key takeaway here. The

next set of optimizations in inference are workload aware. Inference engines become aware of the workload being served and use that information to take KV cache reuse and KV cache orchestration opportunities to the next level. These policies, these examples are academic papers that show value, but you can't really take them to production as is because, you know, each is a fork of LM. Each they're not composable. They're

building on their own. LMD aims to provide the missing layer to allow for creating these policies. And uh thank you all. We're out of time. This will be progressing in LMD. Uh do follow us and if you're interested, contribute as well. Thank you.