[RFC]: Lazy CUDA Graph capture

[lionelvillard]

Motivation.

Currently vLLM captures cudagraphs as part of the engine initialization significantly slowing down vLLM startup time. By default, vLLM captures 67 graphs, which depending on model size and GPU type, can take more than 10s. This is not great UX (see #19824 for details).

In addition, It's most unlikely that all 67 graphs are actually needed, wasting both time and space.

Proposed Change.

We propose to capture cudagraphs lazily. Instead of performing dummy runs during the engine initialization phase, the idea is to do those runs somewhere in the CUDA piecewise backend, and only for the current runtime shape if not cached already.

Exact implementation needs to be worked out.

Feedback Period.

one week

CC List.

@ProExpertProg @aarnphm @charlesfrye

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[ProExpertProg]

Cc @youkaichao @zou3519

[ProExpertProg]

A few thoughts:

• I think the current implementation in CudaPiecewiseBackend already kind-of supports lazy capture. It's just always initiated from the model runner during capture_model.
• We should decide the interface for this, and make it robust. Whether there are separate eager_cudagraph_sizes and lazy_cudagraph_sizes, or if lazy is a subset of all, we should make sure that 1. specifying this on the CLI is intuitive and 2. that the config is the single source of truth for this.
• We could also try to add a feature to the scheduler where it lazily captures cudagraphs if there are no requests to be scheduled. But I'm not sure how common that is in the use-cases we care about (especially auto-scaling)

[zou3519]

A mode to do lazy cudagraphs capture seems reasonable. I'm curious if we think this should be the default or not. Doing cudagraphs lazily means that model serving time becomes less predictable.

[charlesfrye]

Nice! We and our users run a lot of vLLM services that autoscale, including scaling from 0, so startup time matters to us. When scaling from 0, the latency has to be incurred anyway, so lazy capture helps us spread the load out without compromising on later inferences.

Some details on that below!

When would you care about capturing CUDA graphs but not want to do it eagerly?

IME, graph capture is most important for 8B and smaller models. Intuitively: CUDA kernel launch overhead is ~microsecond scale per kernel, so you're talking maybe an ~millisecond improvement in ITL, which matters most when your ITL is also sub-10ms. I see roughly those latencies for models with ~4-8 GB of weights on short contexts on 1xH100 with vLLM/SGLang/TRT-LLM, cf these benchmark results. Speed-of-light here for ITL is also 2-5ms.

Because the model weights are smaller, these are also the engines that boot fastest. In our system weights are loaded from storage at ~GB/s. I think that's typical for deployments of "moderate" sophistication, because you'd need IB/RocE and a weight server get to TB/s. So that part of the engine load is at the seconds scale, whereas it's at the minutes scale for larger models.

In this setting, we are sensitive to the extra latency of CUAD graph capture and the ITL benefits. Boots are ~45s without it, so ~10s would add 20%, cf this PR/write-up.

Other comments

We should decide the interface for this, and make it robust. Whether there are separate eager_cudagraph_sizes and lazy_cudagraph_sizes, or if lazy is a subset of all, we should make sure that 1. specifying this on the CLI is intuitive and that the config is the single source of truth for this.

+1 for separate eager and lazy CUDA graph sizes. I expect we'd want to capture one or two graphs eagerly (say, batch size one and max) and leave the other, less common, sizes to lazy capture.

I'm curious if we think this should be the default or not. Doing cudagraphs lazily means that model serving time becomes less predictable.

IME developing and iterating on a lot of shorter-lived vLLM instances, a bit of extra variability is a very appealing trade for faster startup. The most important use case I can think of where this is also true is for RL off of rollouts -- training code is often short-lived and involves a lot of iteration by researchers. Once you finalize your training code for a run, you might still leave lazy capture in place. Tail latencies are less important, but throughput (which depends on average ITL and TTFT) matters.

[BoyuanFeng]

One concern is the performance drop after start up. By current eager CUDAGraph capture, everything happens during startup time. We can say for sure it will give consistently good performance during runtime. But for lazy CUDAGraph capture, we will see abrupt performance drop during runtime and it would be hard to know when the performance will be stable. Many people prefer consistent performance during runtime over slightly shorter start up time.

Comparing with lazy cudagraph capture, probably reducing number of cudagraph capture sizes is better, especially if we know what batch sizes we care about. Also, full-model cudagraph instead of piecewise cudagraph can also reduce graph capture size. For piecewise cudagraph, we are recording #layers x #sizes CUDAGraphs, which could be ~5000 CUDAGraphs for LLAMA-3.1-70B.

Question: What would be the batch size you care about? Currently we record 67 sizes, including [8, 16, 24, 32, ..., 512] + [1,2,4].

[yeqcharlotte]

We propose to capture cudagraphs lazily. Instead of performing dummy runs during the engine initialization phase, the idea is to do those runs somewhere in the CUDA piecewise backend, and only for the current runtime shape if not cached already.

Another downside of doing this is CUDA graph takes substantial amount of memory that couldn't be taken into account in memory profiling. We'd rather see OOMs instantly and instruct users to start with a different gpu memory utilization than seeing unpredicted OOMs in the middle of some running processes when workload changes.

I feel finding a good balance around how much to capture would make sense @BoyuanFeng said. Say maybe we can have different modes depends on compile optimization modes, say allowing batch size intervals to be uniform, exponential, or more sparse.

I would also imagine diminishing returns over certain batch sizes. I think we should do some benchmarking around it.

Running batch sizes is really workload specific. Typical chatbot workload for TP8, TP1 sized models of ~2K input typically fall within batch size 16 - 64 in order to satisfy the TPOT requirements.

[aarnphm]

Say maybe we can have different modes depends on compile optimization modes, say allowing batch size intervals to be uniform, exponential, or more sparse.

This sounds a bit too complex. I felt like given that this is workload specific, a guide on informing users what bs works for certain usage (function calling, rag, long context) would be a great starter.
+1 for reducing the batch size. For smaller models, we can potentially enable full cuda graph, whereas with larger models (70B and up) we can start with capped batch size to maximum of 128 would be good enough.

[lionelvillard]

Another downside of doing this is CUDA graph takes substantial amount of memory that couldn't be taken into account in memory profiling.

right but that can be solved by caching the max GPU memory utilization to use for a given configuration, similar to caching pytorch compilation artifacts. At a high level it seems this could work.

This makes me think we could also cache the result of peak memory usage profiling.

[angkywilliam]

One additional use case I’d like to add is related to LoRA. I’ve noticed that initiating vLLM with a higher max_loras value (36 in my case) significantly increases the startup time to as high as 7 minutes for a 8B model. My hypothesis is that vLLM creates CUDA graph for each LoRA configuration, which leads to a large number of graphs being generated.

It would be ideal if we could combine prebuilding a subset of CUDA graphs with lazy loading. I believe this would be helpful for both merged and LoRA model use cases, allowing commonly used CUDA graphs to be prebuilt while less common configurations to be built on the fly. For the LoRA case in specific, it would be ideal to prebuilt graphs for up to, say, 8 LoRAs and have the remaining LoRA built on the fly as it's a less common traffic pattern.

[qthequartermasterman]

I think having prompt embeddings enabled vs disabled is relevant to this conversation. They're currently only supported in v0 (haven't yet been ported to v1, I'm hoping to do that before the next release 👀 ), but in the v0 engine, they're currently disabled by default partly because it doubled the cuda graph capture time (the graph had to be compiled separately for prompt token id inputs and prompt embedding inputs). Depending on how the CUDA graph compilation occurs in v1, doing this lazily may have consequences for how prompt embeddings should be enabled.