[Core][CUDA Graph] add output buffer for cudagraph to reduce memory footprint

[youkaichao]

Currently, each cudagraph for one batch size will incur additional memory cost of TP * BATCH_SIZE * HIDDEN_SIZE . Then the total memory footprint for all cudagraphs are TP * SUM(BATCH_SIZES) * HIDDEN_SIZE.

If we create output buffer too, the total memory footprint for all cudagraphs are TP * MAX(BATCH_SIZES) * HIDDEN_SIZE, which is constant, regardless of how many graphs we capture.

[youkaichao]

Performance-wise, this adds one memory copy operation into cudagraph. This should be negligible, as the operation will also be recorded in the cudagraph, and the copy size is quite small.

I run the benchmark, and find the throughput diff is within run-to-run variations.

Before this PR:
Throughput: 18.57 requests/s, 9510.07 tokens/s

After this PR:
Throughput: 18.59 requests/s, 9516.03 tokens/s

[sfc-gh-hazhang]

I observe when enabling cuda_graph, the pp branch will use significantly more memory. This issue might be related?

[youkaichao]

I observe when enabling cuda_graph, the pp branch will use significantly more memory. This issue might be related?

Yes it is possible. To cleverly use cudagraph with multiple sizes, we need some technique like sharing buffers, just as what we do for inputs, and what this PR does for outputs.

[youkaichao]

If this PR is merged, more cuda graphs can be captured at no additional memory cost, which leads to further questions:

• should we capture as many cuda graphs as we can (all batch sizes up until the maximum), so that we don't need padding anymore?
• should we reduce _BATCH_SIZE_ALIGNMENT from 8 to 4 or 2 , so that less computation is wasted on padding?
• should we align the batchsize used for tuning moe kernels with the cudagraph?

vllm/benchmarks/kernels/benchmark_mixtral_moe.py

Lines 19 to 22 in fbdb7b3

	for bs in [
	1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 256, 512, 1024, 1536,
	2048, 3072, 4096
	]:

cc @WoosukKwon for cudagraph
cc @pcmoritz for moe kernels

Update: I realize that cudagraph will cost some memory anyway . Therefore we don't need to actively capture many cudagraphs.

[youkaichao]

For llama2-7b:

before this PR:

before capture graph with batch size 256 Used Memory: 19453.89 MB
after capture graph with batch size 256 Used Memory: 19464.02 MB
before capture graph with batch size 248 Used Memory: 19464.02 MB
after capture graph with batch size 248 Used Memory: 19465.95 MB
before capture graph with batch size 1 Used Memory: 19498.69 MB
after capture graph with batch size 1 Used Memory: 19498.70 MB
INFO 05-27 23:21:09 model_runner.py:920] Graph capturing finished in 6 secs.

the first graph takes 11MB memory, the second takes 2MB memory, and in total cudagraphs take 45MB memory.

after this PR:

before capture graph with batch size 256 Used Memory: 19455.89 MB
after capture graph with batch size 256 Used Memory: 19464.02 MB
before capture graph with batch size 248 Used Memory: 19464.02 MB
after capture graph with batch size 248 Used Memory: 19464.02 MB
...
before capture graph with batch size 1 Used Memory: 19464.05 MB
after capture graph with batch size 1 Used Memory: 19464.05 MB
INFO 05-27 23:18:21 model_runner.py:927] Graph capturing finished in 8 secs.

the first graph takes 8MB memory, the second takes 0MB memory, and in total cudagraphs take 8MB memory. We can see the memory used after capturing all graphs (19464.05 MB) is almost the same as just capturing one graph previously (19464.02 MB).

[WoosukKwon]

@rkooo567 I can take this PR if you are ok with that. Just please let me know!

[rkooo567]

oh yeah please go ahead!

[WoosukKwon]

@youkaichao Thanks for the PR! This is a good finding. Please check my comments.

[WoosukKwon]

LGTM!

[WoosukKwon]

@youkaichao Is the CI failure related to this PR?

[youkaichao]

@youkaichao Is the CI failure related to this PR?

investigating.