Lightning Talk: Beyond Generic Spans: Distributed Tracing for Actio... Sally O'Malley & Greg Pereira

This talk covers the integration of tracing into the LMD framework, a Kubernetes-native distributed tracing tool designed for model serving. The speaker explains their approach to adding tracing while minimizing code changes, ultimately opting for extensive manual instrumentation of individual components rather than reliance on auto instrumentation. Key topics include the gateway API inference extension initiated by Google, the role of an inference scheduler, and various plugins that enhance decision-making capabilities within LMD. The speaker also discusses the significance of tracing in scenarios such as prefix cache-aware routing and pre-fill decode disaggregation, illustrating how these methods improve performance and debugging for AI workloads. They conclude by providing insights into enabling tracing in VLLM and LMD with simple configuration steps.

I'm Sally. I am going to talk about tracing VL and L M D. I added tracing to L M D over the past few months or even longer. It took a while. So what is L M D? It's a Kubernetes native distributed tracing framework. It sort of sits on top of VL or some other model serving engine. It It adds some advanced routing and scaling capabilities. And

so when I when I started out to add tracing in L M D, I wanted to do it with as minimal code changes as possible. So I submitted a proposal to just use auto instrumentation. And I also tried to use the auto EBPF instrumentation and I quickly found that with auto it it it doesn't give the insights that you need with model serving. Things like, you know,

KV cache state. There are a lot of decisions to be made with regards to routing and I wanted to get that full picture. So I went down the route of instrumenting each individual component. The gateway API inference extension is is a project initiated by Google and Cube Sig. It's It's a X proc gateway that adds an endpoint picker protocol that is kind of the scoring engine for

all these decisions that I'll talk about in the next slides. And so that was there but L M D added an inference scheduler which imports all those gateway libraries and the EPP but it adds on top of EPP some plugins. I think currently there are like 16 different plugins or profiles or filters that you can add and that is what makes LLMD the decision maker. And it

also has a KV cache index score and then if you're using pre-filled decode disaggregation, there's a sidecar that gets involved to handle the the transfer of the KV cache and and the decision to go to pre-fill or decode. And then finally you get to VLLM. VLLM tracing was already in the code base, but when I started it was only in V1 and they hadn't ported it to

V2. So over like the 6 months of doing this, I had to wait I didn't do it myself. VLLM brought back in tracing into V2, which is what LLMD uses. And this was a lot, but it's in as of like the the last release. You can imagine how hard it was to maintain all of these working branches of all of these repositories with all of the different

images that had tracing enabled. It was quite a it was quite a journey, but yes, I am going to go through a couple of scenarios where tracing really makes sense for your model serving and one of those is prefix cache aware routing. with LLMD, you can uh get the benefit of a faster time to first token if you route this is just a static picture. I'm going

to show you the live one at the end. If you route your request to a warm KV cache VLLM endpoint. Uh so that's exactly what happens when you enable this EPP plugin. The um scheduler, the inference scheduler that I showed in the previous slide will uh uses the VLLM KV events to get a picture of the KV cache for each VLLM endpoint that you may have. And

uh it uses block level uh cache hit ratio uh as a signal for like the prefix locality. By prefix, I mean the prompt that is you know, in the request. And so uh what happens is between the EPP and the KV index, the um the VLLM pod is chosen with the warmest KV cache, and you get the benefit of a a fast time to first token. Uh

I can show this live at the end again. But here's a here's a picture of an actual trace um from precise prefix cache aware where all of the uh spans were lit up. So, the spans I added were uh the pref the in the EPP pod, the the um the scoring that happens to find how many blocks had your um prompt chunks with you know, and then

um the KV cache has to like you have to look up the KV cache storage to find exactly uh um the hits, and then uh it shows a VLLM. So, that's an actual request, and you can uh um an actual trace, and you can expand each of those spans, uh and you can expand them one more level that I don't have here, but you get all of

the you get the whole story of what happened with a single request, uh the decisions that were made, the the the possible candidates that were there. Uh and so, this is very useful when you're just debugging, maybe um trying find the characteristics of your of your AI workload to see which of these plugins for LMD make the most sense. Uh so, that's the precise prefix cache aware

routing. Uh another scenario when serving it would be um perhaps you want to enable pre- pre-fill decode disaggregation. And that's slightly different. Uh pre-fill decode disaggregation makes when you have uh a lot of prompts dynamic prompts that maybe aren't in the KV cache. And uh very long prompts. And so, sometimes it makes sense to uh separate the pre-fill, which is the part of VLLM that processes the

prompt, from decode, which is the part of VLLM that uh generates the tokens. So, the processing of the prompts is very compute heavy. And so, with very long prompts, it makes sense to separate that into a separate VLLM and then have a super fast kind of connector uh using RDMA or um NicStl is the one that is usually used with LMD. And so, um that's This is

what that looks like. The uh the request comes in and then that gateway and the EPP, that's all in the inference scheduler component that I showed you in that slide. Uh the it decides through the EPP kind of decision maker and scoring it decides whether it makes sense to send your prompt uh to pre-fill. And if it does, then the PD proxy gets involved to transfer your

request over to pre-fill and then the KV transfer very quickly um transfers it over. And if the prompt processing is expensive enough to outweigh the um the cost of the overhead of all of those transfers, that's when PD makes sense. It doesn't make sense always, but with tracing you can track all these things and figure out for your workload if it makes sense. And here's what an

actual trace looks like. With PD, you have different spans that are lit up. Things like did PD happen and is it a yes or no? So you that is that's prefill decode disaggregation. Now VLM has one span. It's the LLM request. Actually has a lot of spans in the very startup phase, um the LLM request is the one that you'll see in all of your all of

your traces and this has all of those GenAI semantic conventions that you're used to seeing maybe in any tracing with any AI workload. Things like time to first token or prompt tokens which model are you using? Don't take a picture of me when I'm drinking water. Just kidding. Thank you. All right, so and with VL you can enable tracing in VLM without LMD. It's just a flag

down there among all of the 400 engine hours of VLM. If you find that OTLP traces endpoint, you're good to go and if you have a collector backend. LMD is very complicated to set up. It's very um so I I wanted to make tracing very easy. So I worked with obviously wasn't just me. I keep saying me, but the person that was supposed to be next to

me, Greg Pereira, is amazing. He really helped get all of the tracing over the finish line as well as Maroon. Maybe he's back there, but anyways, to enable tracing you just configure a couple of flags in our helm charts, and uh uh as as as com- as complicated as it is to to get it uh to get the instrumentation, and it's very easy to enable. And I

wanted to mention these well-lit paths or the guides that you can find in LLMD. Uh that's where I got the two scenario guides uh that I showed kind of showed today. Uh there are many many others, so check out the guides in LLMD. Oop, 9 seconds left. Walk away with this. There are two flags to enable. Uh you get end-to-end visibility uh from the gateway to the

EPP to the KV cache. Uh to the PD proxy sometimes to vLLM. It's a full story. It will help you uh characterize your AI workloads. And uh and this is how enterprise model serving is going to work is with um a whole framework around model serving. I think that that is it. I could show you my live web app that Claude made for me. It's so cool.

So, I'll just show it as as we finish out here. So, the the decisions are being made in the EPP and KV cache that vLLM is chosen, and uh here's the prefill decode. The request comes in. The um inference scheduler and EPP decides whether to route to prefill or decode. Uh when it makes sense, it routes to prefill. When it doesn't, it goes straight to decode. And

uh with that fast KV transfer, uh you get a um you get the time to first token ben- uh benefit. So, that's it. Thanks.