[Core][5/N] Fully working chunked prefill e2e #3884

rkooo567 · 2024-04-06T13:14:34Z

This PR is a part of the RFC #3130.

This PR enables chunked prefill e2e. Note that chunked prefill is an experimental feature now (though it is actively used within Anyscale), and I will start serious benchmark with this PR.

The feature can be enabled by using enable_chunked_prefill. The chunking is done based on max_num_batched_tokens (a.k.a token budget). This is the same way as described in SARATHI paper https://arxiv.org/abs/2308.16369. It also means we can have up to maximum 2 chunked prefill at any given time.

This PR

Allow to put 2 attention metadata, one for prefill and one for decode.
- it is done by broadcasting metadata twice. It can be more optimized by coelescing tensors better, but I made this way for simplicity. The normal path should just use 1 broadcast as usual.
Allow to run attention backend when there are prefill and decode mixed up.
Ignore the generated token if chunked prefill is enabled.
Add a chunked prefill test with various chunk size, cuda graph, and tp settings.

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!

rkooo567 · 2024-04-06T16:21:14Z

vllm/config.py

-            # higher throughput.
-            self.max_num_batched_tokens = max(max_model_len, 2048)
+            if enable_chunked_prefill:
+                # For chunked prefill, choose the well-tuned batch size.


This is arbitrary value. We should profile again.

I would imagine this is hardware/model dependent. @AgrawalAmey any suggestion for a good default value?

actually, there is a no good default -- all depends on application requirements, model, hw, etc. I have an experimental auto-tuner which sets the maximum chunk size that can satisfy the application requirements. It is pretty straight forward, it just takes tbt target and find out the max token budget that can satisfy it.

vllm/core/scheduler.py

rkooo567 · 2024-04-06T16:31:55Z

vllm/attention/backends/torch_sdpa.py

@@ -49,17 +50,14 @@ def copy_blocks(


 @dataclass
-class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
+class TorchSDPAMetadata(AttentionMetadataPerStage, PagedAttentionMetadata):


QQ: Currently, do we run all attention backend in tests?

PR welcomed!

I will create a PR for it

rkooo567 · 2024-04-06T16:35:02Z

tests/worker/test_model_runner.py

+    assert len(return_prompt_lens) == 0
+
+
+# SANG-TODO Test chunked prefill case.


This one is WIP

rkooo567 · 2024-04-06T17:32:45Z

Initial throughput benchmark on 7B, 1 GPU

# not enabled
Throughput: 1.98 requests/s, 974.93 tokens/s
# enabled
Throughput: 2.02 requests/s, 995.76 tokens/s

I am going to start further benchmark at better machines and higher TP next week

scv119 · 2024-04-06T20:47:28Z

for the benchmarks, does it enable cuda-graph? also worth noting the detail of benchmarks (i.e. type of GPU and benchmark setup)

ywang96 · 2024-04-07T04:30:01Z

@rkooo567 If this PR is ready, I'm happy to test it with some benchmarking as well!

rkooo567 · 2024-04-07T12:53:34Z

@ywang96 I think it may actually work. The test failures now may be just test specific issues. The basic_correctness_test seems pass when I ran it (including cuda graph + tp > 1).

@scv119 actually the original RFC doesn't include cuda graph yet after talking with Woosuk (he mentioned he didn't see much improvement in cuda graph if batch size > 256). I will run our internal with decode-only cuda graph (oss status quo) vs whole cuda graph to see the actual impact and do a follow up. Have you observed big perf difference here?

rkooo567 · 2024-04-07T13:18:00Z

also want to emphaisze the benchmark above runs 1 A10 GPU + llama 7B. It was just a dry run. I will do more serious benchmark next week, and we may need cuda graph for chunked prefill to (now it is enabled only when batch only contains decode)

scv119 · 2024-04-07T20:34:42Z

actually the original RFC doesn't include cuda graph yet after talking with Woosuk (he mentioned he didn't see much improvement in cuda graph if batch size > 256). I will run our internal with decode-only cuda graph (oss status quo) vs whole cuda graph to see the actual impact and do a follow up. Have you observed big perf difference here

my sense is cuda-graph will help with tp > 1 cases regardless of chunk-prefill enabled or not; but better just run some benchmarks.

simon-mo · 2024-04-08T20:12:58Z

vllm/attention/backends/torch_sdpa.py

@@ -49,17 +50,14 @@ def copy_blocks(


 @dataclass
-class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
+class TorchSDPAMetadata(AttentionMetadataPerStage, PagedAttentionMetadata):


PR welcomed!

simon-mo · 2024-04-08T20:17:49Z

vllm/attention/backends/xformers.py

+        if num_prefill_tokens > 0:
+            prefill_meta = attn_metadata.prefill_metadata
+            assert prefill_meta is not None


Should we the invariant check in the AttentionMetadata data class?

Suggested change

if num_prefill_tokens > 0:

prefill_meta = attn_metadata.prefill_metadata

assert prefill_meta is not None

if prefill_meta: = attn_metadata.prefill_metadata:

assert prefill_meta is not None

Python 3.8's valrus operator is very useful here. But it does need num_prefill_tokens invariant to be verified somewhere else?

moved invariant check and used warlus operator

I'm not a fan of Walrus operator in python and I believe it's pretty controversial. I would vote for limiting it's use in vLLM. @simon-mo what do you think?

But this is a small thing. I don't think we should block any PR with this.

simon-mo · 2024-04-08T20:19:01Z

vllm/attention/backends/abstract.py

@WoosukKwon @zhuohan123 would love your feedback on this interface change where the new AttentionMetadata can contains both prefill and decode stage metadata. I think this is a necessary change but would like to hear your thought on the interface design.

simon-mo · 2024-04-08T20:20:29Z

vllm/attention/backends/abstract.py

+    slot_mapping: torch.Tensor
+    kv_cache_dtype: str


group these two separately from the prefill/decode. i guess it make sense here as a general arguments instead of the per stage ones.

simon-mo · 2024-04-08T20:21:41Z

vllm/config.py

-            # higher throughput.
-            self.max_num_batched_tokens = max(max_model_len, 2048)
+            if enable_chunked_prefill:
+                # For chunked prefill, choose the well-tuned batch size.


I would imagine this is hardware/model dependent. @AgrawalAmey any suggestion for a good default value?

vllm/config.py

vllm/core/scheduler.py

vllm/worker/model_runner.py

vllm/attention/backends/xformers.py

AgrawalAmey · 2024-04-10T03:11:40Z

vllm/attention/backends/xformers.py

-                    attn_metadata.prompt_lens_tensor,
-                    attn_metadata.context_lens,
-                    attn_metadata.max_subquery_len,
+                    prefill_meta.block_tables,


@rkooo567 you mentioned this kernel is really slow right? should we show some warning if someone uses this backend with chunked prefills?

Hmm feel like that's a bit weird because there's no alternative option. I think we should probably just mark it as experimental feature for now?

I can also do some additional e2e benchmark. When I ran it with 7B, the perf was actually not very different

AgrawalAmey · 2024-04-10T03:14:02Z

vllm/attention/backends/flash_attn.py

            else:
                # prefix-enabled attention
                # TODO(Hai) this triton kernel has regression issue (broke) to
                # deal with different data types between KV and FP8 KV cache,
                # to be addressed separately.
-                output = PagedAttention.forward_prefix(
+                output[:num_prefill_tokens] = PagedAttention.forward_prefix(


with flash_attn_varlen_func, we can do both prefill and decode attention together right?

Yes! I heard that it doesn't make much perf change, but we can iterate on this once the API supports larger page size

rkooo567 · 2024-04-10T07:34:12Z

vllm/config.py

+                # for higher throughput.
+                self.max_num_batched_tokens = max(max_model_len, 2048)
+        if enable_chunked_prefill:
+            logger.info("Chunked prefill is enabled (EXPERIMENTAL).")


cc @AgrawalAmey

rkooo567 · 2024-04-10T07:57:49Z

All comments are addressed, and I believe the tests are passing now

zhuohan123

@rkooo567 Let's chat offline about this PR? My high-level comment: with chunked prefill, we should no longer distinguish prefill and decode outside the attention layer, which should simplify the code significantly.

vllm/attention/backends/abstract.py

zhuohan123 · 2024-04-10T08:05:06Z

vllm/attention/backends/torch_sdpa.py

+        assert query.shape[0] == num_prefill_tokens
+        assert decode_query.shape[0] == num_decode_tokens
+
+        if prefill_meta := attn_metadata.prefill_metadata:


Small Python style question: is := widely used now? I personally prefer the more old-school way.

This was based on simon's feedback. #3884 (comment). Lmk if you want to just revert this

zhuohan123 · 2024-04-10T08:14:50Z

vllm/model_executor/parallel_utils/communication_op.py

+        async_handles = []
        for key, value in metadata_list:
            if isinstance(value, TensorMetadata):
                tensor = tensor_dict[key]
-                torch.distributed.broadcast(tensor, src=src, group=group)
+                async_handles.append(
+                    torch.distributed.broadcast(tensor,
+                                                src=src,
+                                                group=group,
+                                                async_op=True))
+        for async_handle in async_handles:
+            async_handle.wait()
+


Thanks for optimizing this!

zhuohan123 · 2024-04-10T08:15:39Z

vllm/worker/model_runner.py

+class PreparePromptMetadata(NamedTuple):
+    input_tokens: List[int]
+    input_positions: List[int]
+    attn_metadata: Optional[AttentionMetadataPerStage]
+    prompt_lens: List[int]
+    subquery_lens: List[int]
+    lora_index_mapping: List[int]
+    lora_prompt_mapping: List[int]
+    lora_requests: Set[LoRARequest]
+    multi_modal_input: Optional[torch.Tensor]
+    slot_mapping: List[int]
+
+
+class PrepareDecodeMetadata(NamedTuple):
+    input_tokens: List[int]
+    input_positions: List[int]
+    attn_metadata: Optional[AttentionMetadata]
+    lora_index_mapping: List[int]
+    lora_prompt_mapping: List[int]
+    lora_requests: Set[LoRARequest]
+    slot_mapping: List[int]


How is chunked prefill handled here?

This addresses #3884 (comment) comment. It is combined inside prepare_input_tensors

vllm/worker/model_runner.py

rkooo567 · 2024-04-10T10:03:32Z

@zhuohan123 I believe the current style is the best iterative step. (it is also the way it is done in our internal repo). I'd love to do unification as a followup, but I feel like we should do it step by step instead of all in the current PR;

Some of internal metadata is not consistent between decode/prefill. For example, context_len. We should fix these tech debt.
This requires to refactor existing code to be prefill/decode agonistic because chunked prefill co-exist with existing code at least for some time.

Let's talk more in details tmrw!

rkooo567 · 2024-04-10T23:47:33Z

Discussed offline.

We agreed on what's "clean" state. It requires some thoughts, and I will do it as a follow up
The sequence is logprob fix -> cuda graph enablement for prefill -> refactor this part to be clean

jikunshang · 2024-04-11T10:03:42Z

vllm/attention/backends/torch_sdpa.py

                attn_metadata.kv_cache_dtype,
                self.num_kv_heads,
                self.scale,
                self.alibi_slopes,
                kv_scale,
            )
+            assert out.shape == output[num_prefill_tokens:].shape
+            output[num_prefill_tokens:]


should be output[num_prefill_tokens:] = out here?

oops that;s correct. I don't know how the test passes...

oh it is handled 8afca50

We didn't enable much test for cpu path, we will try to add step by step later. thanks for your confirming!:)

casper-hansen · 2024-04-11T16:31:38Z

It also means we can have up to maximum 2 chunked prefill at any given time.

Can we optimize this to allow many more chunks at a time? I thought the main optimization of DeepSpeed-MII was the chunked prefilling, which greatly speeds up the context processing and therefore the throughput.

rkooo567 · 2024-04-12T00:22:07Z

@casper-hansen that's totally doable. Do you happen to know the perf benefits it can bring or a reference sections in a paper that explaisn it? I followed the policy from https://arxiv.org/abs/2308.16369. Also cc @AgrawalAmey for thoughts

casper-hansen · 2024-04-12T17:10:28Z

@rkooo567 Please see this blog from Microsoft.

https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fastgen#3-dynamic-splitfuse-a-novel-prompt-and-generation-composition-strategy

ip

61b7294

rkooo567 changed the title ~~[WIP] Chunked prefill e2e oss~~ [WIP] Chunked prefill e2e Apr 6, 2024

rkooo567 added 3 commits April 6, 2024 09:11

ip

0272344

working e2e

01e5b3b

made it work with flash attn

502bd19

rkooo567 changed the title ~~[WIP] Chunked prefill e2e~~ [Core][5/N] Fully working chunked prefill e2e Apr 6, 2024

rkooo567 mentioned this pull request Apr 6, 2024

[RFC] Upstream Chunked Prefill #3130

Closed

rkooo567 assigned simon-mo Apr 6, 2024

rkooo567 commented Apr 6, 2024

View reviewed changes

fix cpu tests

afa247e

simon-mo mentioned this pull request Apr 6, 2024

[Bug]: Crash using --enable-chunked-prefill #3887

Closed

rkooo567 mentioned this pull request Apr 7, 2024

[WIP][1/N] Chunked Prefill #3106

Closed

simon-mo reviewed Apr 8, 2024

View reviewed changes

zhuohan123 self-assigned this Apr 9, 2024

rkooo567 added 4 commits April 9, 2024 16:36

Merge branch 'main' into chunked-prefill-e2e-done

a4cbe2d

Addressed code review.

e735cc2

addressed other comment

62db33a

fix a test

a18ae3a

AgrawalAmey reviewed Apr 10, 2024

View reviewed changes

vllm/attention/backends/xformers.py Outdated Show resolved Hide resolved

AgrawalAmey reviewed Apr 10, 2024

View reviewed changes

rkooo567 added 3 commits April 9, 2024 22:01

Fixed

4b84904

rocm test fix

5ec4891

Fixed all remaining tests

b814fdb

lint

d01f893

rkooo567 commented Apr 10, 2024

View reviewed changes

zhuohan123 reviewed Apr 10, 2024

View reviewed changes

rkooo567 added 2 commits April 10, 2024 03:25

addressed code review.

346e862

Fixed a broken model test

addf88e

Merge branch 'main' into chunked-prefill-e2e-done

26bfcc3

simon-mo approved these changes Apr 11, 2024

View reviewed changes

simon-mo merged commit 67b4221 into vllm-project:main Apr 11, 2024
35 checks passed

cadedaniel mentioned this pull request Apr 11, 2024

[Speculative decoding 6/9] Integrate speculative decoding with LLMEngine #3894

Merged

jikunshang reviewed Apr 11, 2024

View reviewed changes

SageMoore pushed a commit to neuralmagic/nm-vllm that referenced this pull request Apr 11, 2024

[Core][5/N] Fully working chunked prefill e2e (vllm-project#3884)

09e9e8a

bigPYJ1151 mentioned this pull request Apr 11, 2024

[Hardware][Intel] Isolate CPUModelRunner and ModelRunner for better maintenance #3824

Merged

andy-neuma pushed a commit to neuralmagic/nm-vllm that referenced this pull request Apr 12, 2024

[Core][5/N] Fully working chunked prefill e2e (vllm-project#3884)

ffb2afc

z103cb pushed a commit to z103cb/opendatahub_vllm that referenced this pull request Apr 22, 2024

[Core][5/N] Fully working chunked prefill e2e (vllm-project#3884)

9a4406a

dtrifiro mentioned this pull request May 15, 2024

bump ubi base image tag opendatahub-io/vllm#24

Merged

Temirulan pushed a commit to Temirulan/vllm-whisper that referenced this pull request Sep 6, 2024

[Core][5/N] Fully working chunked prefill e2e (vllm-project#3884)

cd28414

		assert len(return_prompt_lens) == 0


		# SANG-TODO Test chunked prefill case.

[Core][5/N] Fully working chunked prefill e2e #3884

[Core][5/N] Fully working chunked prefill e2e #3884

Conversation

rkooo567 commented Apr 6, 2024 • edited Loading

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

Choose a reason for hiding this comment

rkooo567 commented Apr 6, 2024

scv119 commented Apr 6, 2024

ywang96 commented Apr 7, 2024 • edited Loading

rkooo567 commented Apr 7, 2024

rkooo567 commented Apr 7, 2024

scv119 commented Apr 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhuohan123 Apr 10, 2024 • edited Loading

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 commented Apr 10, 2024

zhuohan123 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

rkooo567 commented Apr 10, 2024 • edited Loading

rkooo567 commented Apr 10, 2024

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

casper-hansen commented Apr 11, 2024

rkooo567 commented Apr 12, 2024 • edited Loading

casper-hansen commented Apr 12, 2024

rkooo567 commented Apr 6, 2024 •

edited

Loading

ywang96 commented Apr 7, 2024 •

edited

Loading

zhuohan123 Apr 10, 2024 •

edited

Loading

rkooo567 commented Apr 10, 2024 •

edited

Loading

rkooo567 commented Apr 12, 2024 •

edited

Loading