[Core] Performance optimization for swap_blocks by cuda kernels

[ApostaC]

This PR is part of the big CPU offloading PR #10874 -- this PR contains the new CUDA kernel implementation for swap_blocks.

Performance benchmark

The numbers are collected on A100-40GB-SXM GPUs

# of pages	New implementation	Old implementation
Swap in 1 page	1.21 GB / second	1.21 GB / second
Swap in 1250 pages	12.6 GB / second	3.11 GB / second
Swap out 1 page	1.22 GB / second	1.21 GB / second
Swap out 1250 pages	12.5 GB / second	3.12 GB / seocond

Notes: CUDA graph compatibility

Currently, it pre-allocates a CPU pin-memory tensor for the blocks_to_swap_in and blocks_to_swap_out. It could support CUDA graphs in the future since the address of the pre-allocated buffer won't change.

I did not include it in the PR currently, and as a next step, I can create a new PR for CUDA graphs.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

some systems do not have pin memory (notably, WSL). we need to take care of that. otherwise this PR LGTM.

[youkaichao]

WSL does not support UVA, either. You can use is_pin_memory_available to determine if this optimization can be used.

[ApostaC]

Sound good, just pushed a new PR fixing this

[Yang-x-Zhao]

Impressive results! I am interested in the benchmark methods of this kernel. I have tried to replicate the results using test_cache.py, but I cannot see similar difference. How to achieve 1250 pages when benchmarking?

Also, what does "pages" refer to in the context of vLLM? It seems to me that you are counting on number of "blocks" in worker.py but not "pages"...

[Uxito-Ada]

Impressive results! I am interested in the benchmark methods of this kernel. I have tried to replicate the results using test_cache.py, but I cannot see similar difference. How to achieve 1250 pages when benchmarking?

Also, what does "pages" refer to in the context of vLLM? It seems to me that you are counting on number of "blocks" in worker.py but not "pages"...

Hi @ApostaC , I have the same question and is it available to share the benchmark in detail?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ApostaC.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[hmellor]

After speaking to @ApostaC I'm closing this as it only optimises the V0 engine. There is work in progress for a similar PR in V1.