[Performance][Core] Optimize the performance of evictor v1 and v2 by applying a priority queue and lazy deletion

[llsj14]

FIX #6923

Summary

• I discovered that the eviction logic with the OrderedDict free_table in Evictor V1 and V2 slows down overall performance (especially TTFT) when using prefix caching mode.
• In some scenarios, utilizing prefix caching mode makes the system slower compared to when prefix caching is not used.
• The evict function is frequently called when allocating a new block, as no block is evicted until the block space is full in prefix caching mode.
• The eviction logic was slow because free_table is declared as an OrderedDict, which is a linked list, and it tries to find a block with content hash (Evictor V1) or block ID (Evictor V2) in this free_table.
• Utilizing a priority queue and lazy deletion helps find the block faster.

Result Verification

• As shown in the following output, the block ID and content hash had the same value between the as-is and to-be states (which is expected).
• With this change, I could make the duration of the evict function much faster.

===============================
evicted_block_id compare:  12010   12010
content_hash_compare:  -7334740008364413937   -7334740008364413937
as-is evict duration:  7.0807114243507385 ms
to-be evict duration:  0.012848526239395142 ms
===============================
evicted_block_id compare:  12038   12038
content_hash_compare:  -7008894356950570757   -7008894356950570757
as-is evict duration:  7.1028973907232285 ms
to-be evict duration:  0.008581206202507019 ms
===============================

Performance

• I checked the TTFT performance using llmperf and the Llama3-8B model with an A100 GPU.
• I benchmarked with 1536 input token length (512 same prefix + 1024 random input) and 512 output token length.
• By applying this commit, I can make the system faster while utilizing prefix caching.
• The speed-up metric is calculated based on the performance without prefix caching mode.

as-is

Model	Num Clients	Block Manager	Prefix Caching	TTFT (mean)	Speed Up
Llama3-8B	16	v2	X	841 ms
Llama3-8B	32	v2	X	1441 ms
Llama3-8B	64	v2	X	2619 ms
Llama3-8B	128	v2	X	4729 ms
Llama3-8B	16	v2	O	1962 ms	0.43 (slowed down)
Llama3-8B	32	v2	O	8382 ms	0.17 (slowed down)
Llama3-8B	64	v2	O	12665 ms	0.21 (slowed down)
Llama3-8B	128	v2	O	22439 ms	0.21 (slowed down)

to-be

Model	Num Clients	Block Manager	Prefix Caching	TTFT (mean)	Speed Up
Llama3-8B	16	v2	O	541 ms	1.55
Llama3-8B	32	v2	O	901 ms	1.60
Llama3-8B	64	v2	O	1563 ms	1.68
Llama3-8B	128	v2	O	2947 ms	1.60

[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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

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

🚀

[youkaichao]

thanks for the contribution!

cc @alexm-neuralmagic @cadedaniel for block manager related optimization.

[Yard1]

that 50 constant should be a defined global.

[llsj14]

@Yard1, thank you for your comments. I have fixed the issue and rebased my code.

[Yard1]

FYI this PR seems to be optimizing the same path #7193

[cadedaniel]

At high level these fixes look great, will need evictor folks to review with more detail (sorry for second ping @robertgshaw2-neuralmagic )

[robertgshaw2-redhat]

At high level these fixes look great, will need evictor folks to review with more detail (sorry for second ping @robertgshaw2-neuralmagic )

Thanks, Alex is going to take a look from out side, since he most recently has been in this codepath optimizing BMv2

[alexm-redhat]

Thanks for revealing this bottleneck and fixing it! It is a good idea to use a heap + dict to quickly access an LRU item. Left some minor comments.

[alexm-redhat]

Nice trick with the -num_hashed_tokens to provide heap sorting.

[alexm-redhat]

Why it was necessary to delay the cleanup? Did you find it to be too slow?

[llsj14]

The reason I applied lazy deletion and event triggered cleanup is that searching specific block and deleting outdated blocks from the heap is O(log n). Thus, I skip and pop outdated blocks by checking the free_table in eviction operation, and only clean up the priority queue when it consumes too much memory with outdated blocks.

Since cleanup itself is O(n log n), calling the cleanup function every time would make the system too slow.

[llsj14]

The ideal scenario is when the cleanup function is not needed, as outdated blocks are naturally popped out during the eviction operation.

[llsj14]

@alexm-neuralmagic, thanks to your comment, I fixed the data type mistake and optimized the performance of the cleanup operation.

I used only the free_table and heapify to create a new priority queue, achieving O(n) complexity.

[alexm-redhat]

Dict is definitely faster here

[alexm-redhat]

I would make this a static class member, since it is used only inside the scope of the class below.

[llsj14]

Thank you, I fixed this

[alexm-redhat]

btw, I would rename the topic of the PR to "[Performance] ....", since it is not a bugfix

[llsj14]

/ready

[llsj14]

I rebased codes to resolve the conflict

[mergify]

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

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

[llsj14]

@alexm-neuralmagic @Yard1
I rebased and tested my code again. I would appreciate your reviews.

[llsj14]

In my local test, the test_eviction_alloc_mixed sometimes passes and sometimes fails.

tests/core/block/test_prefix_caching_block.py ................. [  6%]
............................................................... [ 29%]
............................................................... [ 53%]
............................................................... [ 76%]
............................................................... [100%]
=================== 269 passed, 2 warnings in 6.49s ===================

I believe the assertion in this part is not strictly necessary, because all blocks can be candidates for eviction if they have same last accessed time. The key difference is that the previous code search blocks from the beginning of the free table, while my implementation does not.

@leiwen83 @cadedaniel @comaniac
Could you check whether it would be fine to remove the assertion mentioned above and review my PR please?
-> I just changed my code to make the test pass. I prioritized the block_id to select the earlier one under the same conditions.

[llsj14]

@comaniac
Could you review this PR, please?
This PR was previously reviewed, and I have been testing its stability by running it locally for several months. It has also successfully passed unit tests and CI checks.

[comaniac]

I'm a bit worry about this lazy deletion algorithm as it is pretty hard to understand for others and easy to introduce bugs in corner cases. Here are some possible questions people may ask by reading this code:

1. How a block in the heap not in the free table? A related question is why we need to cleanup the heap.
2. How a block in the heap and the free table could have different last access time?

[llsj14]

@comaniac Thank you for the valuable feedback.
I've added comments regarding the lazy deletion process.

I understand your concerns about the lazy deletion algorithm, as it shows O(n log n) time complexity when triggered. However, since outdated entries are also removed through heap pops, I believe cleanup is not an operation that happens frequently.

In fact, I also considered using doubly linked list and dictionary for this optimization. While these structures are generally O(1), I think that if the key value changes(like num_hashed_tokens in this code) from being solely based on the last accessed time (which always increases), adding entries could then take O(n) time (to make doubly linked list sorted). That’s why I opted for a priority queue... Nevertheless, I acknowledge the concerns about lazy deletion holding outdated entries.

[comaniac]

Yes I used doubly linked list in v1 prefix caching and it works well, but it would be tedious for v0.

[llsj14]

Oh I see. I'll check the v1 implementation later as well.

[comaniac]

Otherwise LGTM