[Core][Hash][Automatic Prefix caching] Accelerating the hashing function by avoiding deep copies #4696

KuntaiDu · 2024-05-08T23:36:25Z

The hashing function implemented in automatic prefix caching is running slowly. To benchmark its performance, I implemented a new benchmarking script benchmarks/benchmark_hashing.py. It shows that the hashing takes 47.98% of the total inference time when all prefix blocks are cached.

After profiling, over 90% of the time is spent on preparing the data for hashing. The reason is that there are multiple deep copies involved while preparing the data.
As shown in the figure,

Line 254, 255, 256 are taking about 30% of the total runtime of function hash_of_block due to deep copy, and the hashing function itself only takes 10% of the hash_of_block's runtime.

To fix this issue, I implemented a new auxiliary function in vllm/sequence.py called get_prefix_token_ids. This function does a similar job as get_token_ids()[0:num_tokens], except for the return value of this function is directly hashable (thus no need to use the deep copies to prepare data for hashing).

After this fix, the hashing function is accelerated by 5x, now hashing only takes 15% of the total inference time.

Side note: this pr assumes that the prompt_token_ids inside SequenceData, once created, will not change (as I will convert prompt_token_ids to tuple and reuse it for all future get_prefix_token_ids request). For now, I will assume that this is true. Please leave comments if you feel this may not hold in certain cases.

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

[Bugfix] for bug fixes.
[CI/Build] for build or continuous integration improvements.
[Doc] for documentation fixes and improvements.
[Model] for adding a new model or improving an existing model. Model name should appear in the title.
[Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
[Kernel] for changes affecting CUDA kernels or other compute kernels.
[Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
[Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
[Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

We adhere to Google Python style guide and Google C++ style guide.
Pass all linter checks. Please use format.sh to format your code.
The code need to be well-documented to ensure future contributors can easily understand the code.
Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

…_ids

…equired

Yard1 · 2024-05-08T23:58:23Z

vllm/sequence.py

@@ -119,6 +119,7 @@ def __init__(
            output_token_ids = []

        self.prompt_token_ids = prompt_token_ids


do we need to store it as both list and a tuple?

Technically we do not have to, as long as there is no code that directly change the value of prompt_token_ids attribute inside the class SequenceGroup.
I did a code search that confirms that there is no such code inside vLLM repo now. The search keyword is ".prompt_token_ids"

And there is no such code that changes prompt_token_ids besides the initialization function.

We should make it a tuple then, IMO. This also signifies prompt tokens are immutable.

However, most of the existing codes are assuming the attribute prompt_token_ids to be of type List[int], and only the hashing function requires it to be Tuple[int] so that it is immutable and thus hashable. From this perspective, I guess it is worthwhile to have an extra copy of prompt_token_ids so that both hashing is fast, and developers can still treat prompt_token_ids as a list.

How about we at least make the attribute with list private and only allow access (and not setting) through a property? Not perfect since it can still be modified in place but better than nothing.

Another idea is to subclass list to create FrozenList with mutable methods set to raise an exception but that's a lot more complex.

Since immutable prompt token IDs is the most important assumption, I'd also suggest to change it to tuple directly. It seems not an issue to me to change all List[int] to Tuple[int, ...]. At least type annotation should not be a blocker.

I have changed the data type of prompt_token_ids to Tuple[int, ...], and fixed typing conflicts caused by this change.

…tion conflict

comaniac

LGTM. Thanks for the timely fix!

…t aligns with the expectation of test cases, plus moving benchmark_hashing.py to a new folder

Yard1 · 2024-05-09T23:08:15Z

vllm/outputs.py

@@ -84,7 +84,7 @@ def __init__(
    ) -> None:
        self.request_id = request_id
        self.prompt = prompt
-        self.prompt_token_ids = tuple(prompt_token_ids)
+        self.prompt_token_ids = list(prompt_token_ids)


do we need to do this? maybe it will be better to change tests instead.

The test case corresponding to this change is:
In tests/conftest.py:

for req_output in req_outputs: ... prompt_ids = req_output.prompt_token_ids ... for sample in req_output.outputs: ... req_sample_output_ids.append(prompt_ids + output_ids) ... ...

And the code prompt_ids + output_ids requires prompt_ids to be list.

I guess processing vllm's output using code like prompt_ids + output_ids may be common in current vllm-based apps. So maybe keeping the prompt_token_ids attribute in RequestOutputs as list would be better for compatibility's sake.

Hmm I think it would be fine to keep it as tuple and cast it to list in the test code tbh

…ert to list

cadedaniel · 2024-05-10T04:44:12Z

btw, we will move to the BlockManagerV2 soon, which has better design for hashing. see

vllm/vllm/core/block/prefix_caching_block.py

Line 546 in e965d46

def content_hash(self) -> Optional[int]:

…rder to be OpenAI-compatible

rkooo567

I feel like we should avoid tuple -> list conversion. I measured locally, and it takes more than 10 us, which is pretty expensive.

rkooo567 · 2024-05-10T12:59:09Z

vllm/outputs.py

@@ -75,7 +75,7 @@ def __init__(
        self,
        request_id: str,
        prompt: str,
-        prompt_token_ids: List[int],
+        prompt_token_ids: Union[List[int], Tuple[int, ...]],


Let's remove tuple input?

rkooo567 · 2024-05-10T12:59:34Z

vllm/sequence.py

@@ -139,7 +139,18 @@ def get_output_len(self) -> int:
        return len(self.output_token_ids)

    def get_token_ids(self) -> List[int]:
-        return self.prompt_token_ids + self.output_token_ids
+        return list(self.prompt_token_ids) + self.output_token_ids


remove tuple input and remove conversion!

Agreed. For performance we need to avoid tuple -> list conversion. How about storing the list version as prompt_token_ids for accessing, and stores the tuple version in _prompt_token_ids_tuple for hashing speedup purposes?

yeah I think that sounds good to me. (it should be fine since prompt tokens are not going to be changed)

Revert this part? (ping me when it is done!)

rkooo567 · 2024-05-10T13:01:08Z

vllm/sequence.py

+    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
+        """Get prefix tokens, and make the return value hashable"""
+        prompt_length = len(self.prompt_token_ids)
+        if num_tokens > prompt_length:


when does this happen?

This happens when calculating hashes for both the user input (i.e. the prompt tokens) and the LLM-generated output (i.e. output tokens ).

Can it happen under normal circumstance or is it only for recomputation case?

Yes it will happen for normal cases (inside function _allocate_last_physical_block in vllm/sequence.py).

rkooo567 · 2024-05-10T13:01:33Z

vllm/sequence.py

        output_token_ids: Optional[List[int]] = None,
    ) -> None:
        if output_token_ids is None:
            output_token_ids = []

-        self.prompt_token_ids = prompt_token_ids
+        self.prompt_token_ids: Tuple[int, ...] = tuple(prompt_token_ids)


Add comment this should not be changed?

Also consider making it private attr

self._prompt_token_id

?

…ompt_token_ids_tuple to speed up hashing

rkooo567

I think you should revert get_token_ids! I will approve it after that

rkooo567 · 2024-05-13T14:36:16Z

vllm/sequence.py

@@ -139,7 +139,18 @@ def get_output_len(self) -> int:
        return len(self.output_token_ids)

    def get_token_ids(self) -> List[int]:
-        return self.prompt_token_ids + self.output_token_ids
+        return list(self.prompt_token_ids) + self.output_token_ids


Revert this part? (ping me when it is done!)

rkooo567 · 2024-05-13T14:37:03Z

vllm/sequence.py

+    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
+        """Get prefix tokens, and make the return value hashable"""
+        prompt_length = len(self.prompt_token_ids)
+        if num_tokens > prompt_length:


Can it happen under normal circumstance or is it only for recomputation case?

KuntaiDu · 2024-05-14T05:06:47Z

I think you should revert get_token_ids! I will approve it after that

Done

rkooo567 · 2024-05-14T12:28:03Z

before this PR: Hashing took 0.16 seconds,14.13% of the total runtime.
after this PR: 4.75%

…ion by avoiding deep copies (vllm-project#4696)

KuntaiDu added 4 commits May 8, 2024 22:19

Kuntai: improve hashing to avoid multiple deep copies on prompt_token…

16fd11a

…_ids

Kuntai: update benchmark_hashing.py so that no huggingface token is r…

e9dc917

…equired

Kuntai: fix formatting issues

8112e0d

Kuntai: fix issues after formatting

57fd939

Yard1 reviewed May 8, 2024

View reviewed changes

Kuntai: adjust prompt_token_ids to tuple, and fix various type annota…

803572e

…tion conflict

KuntaiDu requested review from comaniac and Yard1 May 9, 2024 20:56

comaniac approved these changes May 9, 2024

View reviewed changes

Kuntai: let prompt_token_ids in RequestOutputs be List[int] so that i…

b9b38c0

…t aligns with the expectation of test cases, plus moving benchmark_hashing.py to a new folder

Yard1 reviewed May 9, 2024

View reviewed changes

KuntaiDu added 2 commits May 10, 2024 00:01

Kuntai: update the test case of speculative decoding: explicitly conv…

e0691a6

…ert to list

Kuntai: keep the prompt_ids as Tuple and change the test cases instead

bf498b0

Kuntai: revert RequestOutput.prompt_token_ids back to List[int], in o…

b352246

…rder to be OpenAI-compatible

rkooo567 requested changes May 10, 2024

View reviewed changes

KuntaiDu added 2 commits May 10, 2024 15:13

Kuntai: make prompt_token_ids list, and add new private attribute _pr…

a868884

…ompt_token_ids_tuple to speed up hashing

Kuntai: remove a list conversion

49d994d

KuntaiDu requested review from rkooo567 and Yard1 May 10, 2024 18:01

rkooo567 reviewed May 13, 2024

View reviewed changes

Kuntai: remove list() conversion in get_token_ids

e95d660

KuntaiDu requested a review from rkooo567 May 14, 2024 05:06

rkooo567 approved these changes May 14, 2024

View reviewed changes

rkooo567 merged commit ccb63a8 into vllm-project:main May 14, 2024
59 checks passed

robertgshaw2-neuralmagic pushed a commit to neuralmagic/nm-vllm that referenced this pull request May 19, 2024

[Core][Hash][Automatic Prefix caching] Accelerating the hashing funct…

feed62d

…ion by avoiding deep copies (vllm-project#4696)

KuntaiDu deleted the kuntai-hashing branch May 20, 2024 17:29

dtrifiro pushed a commit to dtrifiro/vllm that referenced this pull request May 21, 2024

[Core][Hash][Automatic Prefix caching] Accelerating the hashing funct…

ed2d743

…ion by avoiding deep copies (vllm-project#4696)

tybalex pushed a commit to tybalex/vllm-function-call that referenced this pull request May 25, 2024

[Core][Hash][Automatic Prefix caching] Accelerating the hashing funct…

a09d8fc

…ion by avoiding deep copies (vllm-project#4696)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core][Hash][Automatic Prefix caching] Accelerating the hashing function by avoiding deep copies #4696

[Core][Hash][Automatic Prefix caching] Accelerating the hashing function by avoiding deep copies #4696

KuntaiDu commented May 8, 2024

Yard1 May 8, 2024

KuntaiDu May 9, 2024

Yard1 May 9, 2024

KuntaiDu May 9, 2024

Yard1 May 9, 2024

comaniac May 9, 2024

KuntaiDu May 9, 2024 •

edited

comaniac left a comment

Yard1 May 9, 2024 •

edited

KuntaiDu May 9, 2024 •

edited

Yard1 May 10, 2024

cadedaniel commented May 10, 2024

rkooo567 left a comment

rkooo567 May 10, 2024

rkooo567 May 10, 2024

KuntaiDu May 10, 2024

rkooo567 May 10, 2024

rkooo567 May 13, 2024 •

edited

rkooo567 May 10, 2024

KuntaiDu May 10, 2024

rkooo567 May 13, 2024

KuntaiDu May 14, 2024 •

edited

rkooo567 May 10, 2024

rkooo567 May 10, 2024

rkooo567 left a comment

rkooo567 May 13, 2024 •

edited

rkooo567 May 13, 2024

KuntaiDu commented May 14, 2024

rkooo567 commented May 14, 2024

		@@ -119,6 +119,7 @@ def __init__(
		output_token_ids = []

		self.prompt_token_ids = prompt_token_ids

[Core][Hash][Automatic Prefix caching] Accelerating the hashing function by avoiding deep copies #4696

[Core][Hash][Automatic Prefix caching] Accelerating the hashing function by avoiding deep copies #4696

Conversation

KuntaiDu commented May 8, 2024

PR Title and Classification

Code Quality

Notes for Large Changes

What to Expect for the Reviews

Thank You

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu May 9, 2024 • edited

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

Yard1 May 9, 2024 • edited

Choose a reason for hiding this comment

KuntaiDu May 9, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cadedaniel commented May 10, 2024

rkooo567 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 May 13, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu May 14, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 left a comment

Choose a reason for hiding this comment

rkooo567 May 13, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu commented May 14, 2024

rkooo567 commented May 14, 2024

KuntaiDu May 9, 2024 •

edited

Yard1 May 9, 2024 •

edited

KuntaiDu May 9, 2024 •

edited

rkooo567 May 13, 2024 •

edited

KuntaiDu May 14, 2024 •

edited

rkooo567 May 13, 2024 •

edited