Lightning Talk: Slash LLM Cold-Start Times by Pre-distributing GPU... Billy McFall & Maryam Tahhan

This talk covers strategies for reducing the startup time of large language model (LLM) deployments in production environments. The speaker discusses the challenges posed by the just-in-time (JIT) compilation process in PyTorch, which can lead to idle GPU time during workload initialization. To address this, Red Hat has developed a utility called Model Cache Vault (MCV), which packages the JIT output in an OCI-compliant image for easier distribution across a Kubernetes cluster. Alongside MCV, the speaker introduces the GPU Kernel Manager (GKM), a Kubernetes operator that verifies the integrity of these images and facilitates the extraction of GPU kernel caches into persistent volume claims. By utilizing these tools, organizations can achieve significant reductions in startup times while maintaining existing workflows with no changes needed to their PyTorch applications.

Hello, my name's Bill McFall. I'm from Red Hat. Um my uh colleague Miriam was supposed to present with me, but she got pulled in to do um one of the keynotes and also doing a demo at the booth, so I told her I'd take over, but Miriam was another key contributor to what we're going to talk about, so I want to give her some credit, but um

she was unable to help me out today. All right. Um All right, so I want to talk to you today a little bit about how to um reduce the startup time of some of your LLM deployments. So, if you're running um LLM in production in um large scales, you're probably paying what we're calling the JIT tax. So, while your PyTorch workload is spinning up, in the background,

PyTorch is doing a JIT compile while your expensive GPUs are sitting there idle and not doing anything waiting for your workload to come up. what we'd like to do is try to figure out a way to reduce this JIT tax, cuz the problem is you're paying this tax over and over again. Every time the pod spins up in the same cluster, you're real rebuilding the exact same

GPU kernel across the cluster across for every node, and even on the same node, if you're running some type of replica set or your node has to restart, you're still rebuilding that same kernel even though it's already been built. So, we want to try to figure out how to fix that. So, in order to do that, um we want to dig a little bit deeper into what's

going on. So, under the covers, PyTorch is calling um torch.compile. Torch.compile is detecting the exact hardware GPU hardware that's running in your system, and then it is performing a just-in-time compile or JIT compile to build these GPU-specific kernels. Now, these kernels while they're extremely they're specialized for the GPU that it detects. So, this JIT is very important cuz you need that specialized hardware. But, you're paying the

tax at startup time. The output of the JIT compile is stored in a local directory. We're calling it the GPU kernel cache. So, if we can capture this cache and distribute it, maybe we'll be able to speed up um our startup times. So, at Red Hat, what we've done is we've created a utility program. We're calling it Model Cache Vault or MCV. So, this utility will package

up this this um GPU kernel cache directory that we've talked about. It'll package up in a standard OCI compliant image. Um because that is a standard OCI image, all your existing tools and utilities will still work. Um things like Docker, Podman, Buildah, whatever you typically use, you can still use cuz it's just a standard container. Also, because we're in production, um security is very important. So, even

though it's optional for what we're doing, we highly recommend running the Sigstore's Cosign to cryptographically sign these OCI images. So, that when you later use them, you can guarantee that they haven't been tampered with, giving you your um supply chain security that you really want in production. So, in addition to the MCV that we've created, we've been working on uh the GPU kernel manager. So, our GKM.

So, GKM is a Kubernetes operator. So, what it does is it pulls down this OCI registry this OCI image from the registry. If you signed it, it'll verify um that the signature is still valid and that the image has not been tampered with. Once it's guaranteed that you have an untampered image. It'll um run the MCV in a Kubernetes job to extract out the contents of the

OCI image into uh Kubernetes persistent volume claim or PVC. And then it'll make sure your PVC's been distributed across the cluster as needed. And then um you can mount volume mount the PVC into your workload. So when the PVC is mounted and your pod comes up, to the pod, it just looks like another directory. So then your JIT compile comes in, it sees the directory is already

in place, does a warm start basically cuz it says, "Oh, I see my directory. I don't need to recompile." and keeps going. So to get the full chain, first you want to do you still need to do the JIT compile at least once. It's important to do that. But you can do this either offline on your training sessions or you can do it on your production cluster

just as long as you have the same GPU architecture and the same input parameters. MCV will then you be used to package up GPU kernel cache directory and then optionally sign it but highly recommended in production to sign it using cosign push up to your registry. GKM will pull it down, verify that it hasn't been tampered with, and extract it out into a PVC on the target

node. And then after that, the pod fires up, sees it's got a warm cache, and um continues on its way. So by just distributing this cache across your cluster, you can save minutes of time. I didn't mention that earlier but some of these um GPU JIT compiles can take minutes. So it's not we're not talking like some of these can get depending on the size can get

really big. Also, this requires zero PyTorch changes. So your workload pod is exactly as it was before. It works um with zero The only thing that potentially we need to change is your pod spec that's running the um workload, but that can be use a web hook to mutate it to add the volume out if you so choose. and also it's fully transparent so that your work

part the code knows no difference. So we've given you a blueprint on how to speed up your production environment. We're currently at Red Hat ET / GKM, but we've been working with K-Serve community and are planning to move the functionality up to there under K-Serve. And I know potentially not everyone in this community uses Kubernetes, but hopefully that the MCV will still be useful to package up

and put manage images outside and then you can just reproduce the same um GKM operator functionality, you know, in your own environment. So thank you for your time. Hope you learned something today and have a good day.