[Bugfix] Enable chunked-prefill and prefix cache with flash-attn backend #6144

sighingnow · 2024-07-04T15:30:18Z

As the flash-attn becomes the default attention backend on NV GPUs (thanks to vllm-flash-attention fork), it is possible to enable chunked prefill with prefix caching at the same time in vLLM.

FIX #xxxx (link existing issues this PR will resolve)

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!

sighingnow · 2024-07-05T01:20:39Z

Hi @rkooo567, could you please take a look at this pull request?

Thanks!

Juelianqvq · 2024-07-05T02:00:08Z

Hi @rkooo567, could you please take a look at this pull request?

Thanks!

I make a try and output consistency cannot be guaranteed under some circumstances.

sighingnow · 2024-07-05T02:37:27Z

I make a try and output consistency cannot be guaranteed under some circumstances.

May I know more about the failed case? How could I reproduce the failure?

Juelianqvq · 2024-07-05T02:48:09Z

May I know more about the failed case? How could I reproduce the failure?

It was not easy to examplify. I test it using Locust with high concurrency and some of the outputs become weird.

rkooo567 · 2024-07-18T06:28:15Z

I will take a look at it tomorrow.

comaniac

Overall LGTM. btw can you add a test case?

comaniac · 2024-07-19T04:47:14Z

vllm/worker/model_runner.py

@@ -277,7 +277,8 @@ def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
            prefix_cache_hit = (computed_block_nums is not None
                                and len(computed_block_nums) > 0
                                and self.sliding_window is None and is_prompt)
-            if self.chunked_prefill_enabled and prefix_cache_hit:
+            if (self.chunked_prefill_enabled and prefix_cache_hit
+                    and self.attn_backend.get_name() != "flash-attn"):
                raise RuntimeError(
                    "chunked prefill cannot be used with prefix caching now.")


Chunked prefill and prefix caching can only be used simultaneously with flash-attn backend

comaniac · 2024-07-19T04:49:44Z

vllm/worker/model_runner.py

+                #       cached parts of the sequence.
+                #   - prefix_cache_len >= seq_len:
+                #     - it means the current partial sequence is fully cache
+                #       hited, and no further computation is needed.


Suggested change

# hited, and no further computation is needed.

# hit, and no further computation is needed.

rkooo567

Looks pretty clean! Can you add unit tests? I think it'd be great to have 1 for worker/test_model_runner.py (just to verify metadata) and 1 e2e test

rkooo567 · 2024-07-19T06:52:12Z

vllm/worker/model_runner.py

@@ -277,7 +277,8 @@ def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
            prefix_cache_hit = (computed_block_nums is not None
                                and len(computed_block_nums) > 0
                                and self.sliding_window is None and is_prompt)
-            if self.chunked_prefill_enabled and prefix_cache_hit:
+            if (self.chunked_prefill_enabled and prefix_cache_hit
+                    and self.attn_backend.get_name() != "flash-attn"):
                raise RuntimeError(
                    "chunked prefill cannot be used with prefix caching now.")


Update the error message that it is only suppored by flash attn backend (also please write how to enable flash attn backend. I.e., VLLM_ATTENTION_BACKEND=FLASH_ATTN)

sighingnow · 2024-07-25T12:21:49Z

Hi @rkooo567 @comaniac, I have addressed the review comments and added a unittest to check the model inputs, could you help to take another look on this PR?

Thanks!

comaniac

Overall LGTM! Just some miner comments for comments and refactoring.

comaniac · 2024-07-25T16:27:05Z

vllm/worker/model_runner.py

+        if (self.chunked_prefill_enabled and prefix_cache_hit
+                and self.attn_backend.get_name() != "flash-attn"):
            raise RuntimeError(
-                "chunked prefill cannot be used with prefix caching now.")
+                "Chunked prefill and prefix caching can only be used "
+                "simultaneously with flash-attn backend, try switching "
+                "to flash-attn backend by setting the environment variable "
+                "\"VLLM_ATTENTION_BACKEND=FLASH_ATTN\"")


prepare_input is now backend agnostic. Can you put this checker to vllm/attention/backends/flashinfer.py and vllm/attention/backends/utils.py? For example in utils.py:

def _add_seq_group( self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup", chunked_prefill_enabled: bool): is_prompt = inter_data.is_prompt block_tables = inter_data.block_tables computed_block_nums = inter_data.computed_block_nums if chunked_prefill_enabled and inter_data.prefix_cache_hit: raise RuntimerError(...)

comaniac · 2024-07-25T16:29:41Z

vllm/worker/model_runner.py

+            # When prefix caching meets chunked prefill, we would be in
+            # the following three conditions:
+            #
+            #   - prefix_cache_len <= context_len:
+            #     - do normal chunked prefill and nothing special
+            #   - context_len < prefix_cache_len < seq_len:
+            #     - advance the context_len to seq_len to perform non-
+            #       cached parts of the sequence.
+            #   - prefix_cache_len >= seq_len:
+            #     - it means the current partial sequence is fully cache
+            #       hit, and no further computation is needed.


nit

Suggested change

# When prefix caching meets chunked prefill, we would be in

# the following three conditions:

#

# - prefix_cache_len <= context_len:

# - do normal chunked prefill and nothing special

# - context_len < prefix_cache_len < seq_len:

# - advance the context_len to seq_len to perform non-

# cached parts of the sequence.

# - prefix_cache_len >= seq_len:

# - it means the current partial sequence is fully cache

# hit, and no further computation is needed.

# When prefix caching meets chunked prefill, we would be in

# one of the following three cases:

# - prefix_cache_len <= context_len:

# Do normal chunked prefill

# - context_len < prefix_cache_len < seq_len:

# Advance the context_len to seq_len to prefill non-

# cached parts of the sequence.

# - prefix_cache_len >= seq_len:

# The current partial sequence is fully cache hit,

# and no further computation is needed.

Another way is embedding the comments to each branch:

if context_len < prefix_cache_len < seq_len: # prefix_cache_len <= context_len: Do normal chunked prefill .... elif seq_len <= prefix_cache_len: # context_len < prefix_cache_len < seq_len: ... ... else: # prefix_cache_len >= seq_len:: ... pass

I'm fine with either way

comaniac · 2024-07-25T16:34:13Z

tests/worker/test_model_runner.py

+
+    seq_lens: List[int] = []
+    seq_group_metadata_list: List[SequenceGroupMetadata] = []
+    block_tables = {0: list(range(128))}


Can you note the block size here for convenience?

comaniac · 2024-07-25T16:34:53Z

tests/worker/test_model_runner.py

+    assert len(input_tokens) == 220
+    assert len(input_positions) == 220


Can you note how 220 was calculated?

sighingnow · 2024-08-06T09:52:20Z

Hi @Juelianqvq could you take another try on this PR to see if previous problems you meet can be reproduced now.

@comaniac Could you please help review changes in this PR again? thanks!

sighingnow · 2024-08-06T14:11:47Z

These two failed test case is not caused by this PR.

Juelianqvq · 2024-08-07T04:21:49Z

Hi @Juelianqvq could you take another try on this PR to see if previous problems you meet can be reproduced now.

@sighingnow Let me test it with our inner dataset consists of high concurreny under differenct scenerios next Monday and give you feedback based on that ASAP. Have you tested with blockmanager v2 yet?

sighingnow · 2024-08-07T04:40:34Z

Hi @Juelianqvq could you take another try on this PR to see if previous problems you meet can be reproduced now.

@sighingnow Let me test it with our inner dataset consists of high concurreny under differenct scenerios next Monday and give you feedback based on that ASAP. Have you tested with blockmanager v2 yet?

Make sure you are working on vllm version with PR #7018 included.

Have you tested with blockmanager v2 yet?

Haven't tested yet, but there should be no difference.

See also: vllm-project#6144

Juelianqvq · 2024-08-12T06:15:27Z

@sighingnow Unfortunately, cudaError is fixed but still can generate garbage sometimes which I've undergone before and made a work-around in #6819. If you don't have such dataset, you can reduce max-num-batched-tokens. Perhaps need to modify on that one in your spare time.

rkooo567 · 2024-08-12T17:50:32Z

@Juelianqvq is it possible to provide a simple unit tests for us?

rkooo567

The PR looks good to me, but the latest comment concerns me if there are still some other issues.

@sighingnow is there any good way to create an e2e test?

sighingnow · 2024-08-13T00:48:43Z

but still can generate garbage sometimes

@Juelianqvq Could you share the command/script that you can reproduce the failure? even if your dataset cannot be available.

is there any good way to create an e2e test?

@rkooo567 An e2e test cannot help the concerns here as we cannot reproduce. The only condition that may lead to incorrect result is hash collision (two non-equal sequences share the same pages). To reproduce such error, it is highly dataset dependent.

Juelianqvq · 2024-08-13T01:44:53Z

@sighingnow Nothing special than add --enable-chunked-prefill -- enable-prefix-caching on finetuned Llama-2-13B.

@sighingnow @rkooo567 Let me clarify what I have done and what problem it is, it is not related to 2-different-sequences but 1-same-sequence which helps me to discover the bug. Given that max-num-batched-tokens set to 512, if I pass a 570 prompt here, it will be chunked into 2 pieces for scheduling, 512 + 58 right? And what will happen if I pass the same request again? Of course, 570 is not fully cached and it will be 511 cached+ 1 recompute + 58 / blocksize * blocksize cached + the rest recompute right? But with this PR, the output becomes shorter than the first round. That's what I met and thought for nearly 2 weeks.

I have not understood the root cause here but I can give you some hints. What confused me is that a sequence being chunked into different scheduling can have changeable cached length, that's the key point and it is related to block manager. Specifically, 570 is split into 512 + 58, and cached length is calculated based on 570(e.g 352) at first and takes effect on 512 part, but after 512 is calculated, cached length of 58(like 384, forget the real number) does not match before (e.g. 352). So I made a work-around to fix the cached length and a little modification on block manager. Hope to find out the real reason pointed out here.

Juelianqvq · 2024-08-13T02:01:10Z

@sighingnow Chunked prefill is connected with chunk size which can be inferred from max-num-batched-tokens, so prompt length greater or less than this value should be included. The reason you have missed to keep at least 1 token for prefill is that you do not select greater prompt length which is always fine in 1-concurrency e2e, but it will cause assertion error in max-query-len > 0 under pressure test.

Moreover, v2-block-manager is another story.

Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>

sighingnow · 2024-08-21T20:43:29Z

@sighingnow Nothing special than add --enable-chunked-prefill -- enable-prefix-caching on finetuned Llama-2-13B.

@sighingnow @rkooo567 Let me clarify what I have done and what problem it is, it is not related to 2-different-sequences but 1-same-sequence which helps me to discover the bug. Given that max-num-batched-tokens set to 512, if I pass a 570 prompt here, it will be chunked into 2 pieces for scheduling, 512 + 58 right? And what will happen if I pass the same request again? Of course, 570 is not fully cached and it will be 511 cached+ 1 recompute + 58 / blocksize * blocksize cached + the rest recompute right? But with this PR, the output becomes shorter than the first round. That's what I met and thought for nearly 2 weeks.

I have not understood the root cause here but I can give you some hints. What confused me is that a sequence being chunked into different scheduling can have changeable cached length, that's the key point and it is related to block manager. Specifically, 570 is split into 512 + 58, and cached length is calculated based on 570(e.g 352) at first and takes effect on 512 part, but after 512 is calculated, cached length of 58(like 384, forget the real number) does not match before (e.g. 352). So I made a work-around to fix the cached length and a little modification on block manager. Hope to find out the real reason pointed out here.

The block manager v1 case has been fixed by 099ab34

sighingnow · 2024-08-21T20:49:08Z

@sighingnow Chunked prefill is connected with chunk size which can be inferred from max-num-batched-tokens, so prompt length greater or less than this value should be included. The reason you have missed to keep at least 1 token for prefill is that you do not select greater prompt length which is always fine in 1-concurrency e2e, but it will cause assertion error in max-query-len > 0 under pressure test.

Moreover, v2-block-manager is another story.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bugfix] Enable chunked-prefill and prefix cache with flash-attn backend #6144

[Bugfix] Enable chunked-prefill and prefix cache with flash-attn backend #6144

sighingnow commented Jul 4, 2024 •

edited

Loading

sighingnow commented Jul 5, 2024

Juelianqvq commented Jul 5, 2024

sighingnow commented Jul 5, 2024

Juelianqvq commented Jul 5, 2024

rkooo567 commented Jul 18, 2024

comaniac left a comment

comaniac Jul 19, 2024

comaniac Jul 19, 2024

rkooo567 left a comment

rkooo567 Jul 19, 2024

sighingnow commented Jul 25, 2024

comaniac left a comment

comaniac Jul 25, 2024

comaniac Jul 25, 2024

comaniac Jul 25, 2024

sighingnow Aug 6, 2024

comaniac Jul 25, 2024

sighingnow Aug 6, 2024

sighingnow commented Aug 6, 2024

sighingnow commented Aug 6, 2024

Juelianqvq commented Aug 7, 2024 •

edited

Loading

sighingnow commented Aug 7, 2024

Juelianqvq commented Aug 12, 2024 •

edited

Loading

rkooo567 commented Aug 12, 2024

rkooo567 left a comment

sighingnow commented Aug 13, 2024

Juelianqvq commented Aug 13, 2024

Juelianqvq commented Aug 13, 2024

sighingnow commented Aug 21, 2024

sighingnow commented Aug 21, 2024

Juelianqvq commented Aug 22, 2024

Juelianqvq commented Aug 23, 2024 •

edited

Loading

sighingnow commented Aug 23, 2024

	# hited, and no further computation is needed.
	# hit, and no further computation is needed.

		assert len(input_tokens) == 220
		assert len(input_positions) == 220

[Bugfix] Enable chunked-prefill and prefix cache with flash-attn backend #6144

[Bugfix] Enable chunked-prefill and prefix cache with flash-attn backend #6144

Conversation

sighingnow commented Jul 4, 2024 • edited Loading

PR Title and Classification

Code Quality

Notes for Large Changes

What to Expect for the Reviews

Thank You

sighingnow commented Jul 5, 2024

Juelianqvq commented Jul 5, 2024

sighingnow commented Jul 5, 2024

Juelianqvq commented Jul 5, 2024

rkooo567 commented Jul 18, 2024

comaniac left a comment

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

Choose a reason for hiding this comment

sighingnow commented Jul 25, 2024

comaniac 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sighingnow commented Aug 6, 2024

sighingnow commented Aug 6, 2024

Juelianqvq commented Aug 7, 2024 • edited Loading

sighingnow commented Aug 7, 2024

Juelianqvq commented Aug 12, 2024 • edited Loading

rkooo567 commented Aug 12, 2024

rkooo567 left a comment

Choose a reason for hiding this comment

sighingnow commented Aug 13, 2024

Juelianqvq commented Aug 13, 2024

Juelianqvq commented Aug 13, 2024

sighingnow commented Aug 21, 2024

sighingnow commented Aug 21, 2024

Juelianqvq commented Aug 22, 2024

Juelianqvq commented Aug 23, 2024 • edited Loading

sighingnow commented Aug 23, 2024

sighingnow commented Jul 4, 2024 •

edited

Loading

Juelianqvq commented Aug 7, 2024 •

edited

Loading

Juelianqvq commented Aug 12, 2024 •

edited

Loading

Juelianqvq commented Aug 23, 2024 •

edited

Loading