[2/N] Chunked prefill data update #3538

rkooo567 · 2024-03-21T01:48:52Z

It is the second PR to enable chunked prefill.

This PR introduces prefill's start and end indexes and API to access them. Also add metadata to SequenceGroupMetadata to indicate if the seq_group is in chunked prefill stage.

Ideally, this mechanism should be unified to computed_num_blocks approach, but I will do this after merging chunked prefill PRs (since it already works e2e).

Related #3130

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!

benchmarks/benchmark_latency.py

vllm/sequence.py

WoosukKwon · 2024-03-27T02:50:56Z

vllm/worker/model_runner.py

@@ -819,8 +840,7 @@ def capture_model(self, kv_caches: List[torch.Tensor]) -> None:
                    context_lens=context_lens[:batch_size],
                    block_tables=block_tables[:batch_size],
                    use_cuda_graph=True,
-                    kv_cache_dtype=self.kv_cache_dtype,
-                )
+                    kv_cache_dtype=self.kv_cache_dtype)


nit: What is this change for?

Hmm not sure why it is formatted. Reverted

rkooo567 · 2024-03-27T02:57:11Z

@WoosukKwon I just pushed a change based on #3538 (comment) comment (the previous commit works, but the commit with this change is not cleaned up yet. I pushed it to demonstrate the difference). Can you take a look at 5e0f87e and see if you feel less hacky this way (or the previous one is better, or lmk if you have the third option I can take).

rkooo567 · 2024-03-27T03:17:02Z

vllm/core/scheduler.py

@@ -27,18 +27,40 @@ class PreemptionMode(enum.Enum):
    RECOMPUTE = enum.auto()


+class ScheduledSequenceGroup:


TODO remove and use tuple

rkooo567 · 2024-03-27T03:18:22Z

vllm/worker/model_runner.py

@@ -819,8 +840,7 @@ def capture_model(self, kv_caches: List[torch.Tensor]) -> None:
                    context_lens=context_lens[:batch_size],
                    block_tables=block_tables[:batch_size],
                    use_cuda_graph=True,
-                    kv_cache_dtype=self.kv_cache_dtype,
-                )
+                    kv_cache_dtype=self.kv_cache_dtype)


Hmm not sure why it is formatted. Reverted

rkooo567

Code is updated based on the initial feedback (to make it more general).

Removed max_chunk_prefill_len. Instead, we will infer this value from max_batched_tokens in the future PR (which is equivalent to the paper https://arxiv.org/pdf/2403.02310.pdf). Thank you @AgrawalAmey !
Instead of updating prefill range, I record "computed_tokens", which is 1 for decoding and chunk size for prefill. Scheduler and SequenceGroupMetadata is updated accordingly.

simon-mo

Thanks for doing this. I left mostly style and clarification comments.

The only thing about the chunk size value during decode. I'm fine with whatever you end up with as long as we have it well documented.

benchmarks/benchmark_latency.py

tests/test_sequence.py

tests/core/test_scheduler.py

vllm/core/scheduler.py

vllm/engine/arg_utils.py

vllm/sequence.py

simon-mo · 2024-03-28T04:18:36Z

vllm/sequence.py

+    def on_recompute(self):
+        """Reset the sequence states for recomputation."""
+        self.data.reset_num_computed_tokens()


maybe call this reset_state so it's easier to understand. recompute is when reset is needed (currently the only scenario but very specific) but this operation is just resetting.

The confusing part (and the reason why I chose more specific term here) is that it actually doesn't "reset" the state because reset sounds like output_token_ids are also cleared. But this API doesn't do that. What about reset_state_for_recompute?

simon-mo · 2024-03-28T04:20:49Z

vllm/sequence.py

+            if is_prompt:
+                self._token_chunk_size = list(seq_data.values())[0].get_len()
+            else:
+                self._token_chunk_size = 1


i guess the value here is not meaningful because it's not accessed for decode, now i wonder whether it make sense to use None

No strong opinion. I just thought it makes more sense as API (that it returns 1 than None). Also less confusing because scheduler returns token_chunk_size == 1 when it is decoding.

rkooo567 · 2024-03-28T04:27:57Z

Thanks for reviewing the PR at late night @simon-mo ! Addressing comments rn...

rkooo567 · 2024-03-28T04:42:25Z

comments are addressed

rkooo567 added 30 commits February 27, 2024 22:55

[1/n] Support efficient reshape caching.

06fe872

[2/n] support flash attention kernel

9a0b6be

oss flash attention works

6947167

in progress

4769a26

flash attn enabled.

963db44

ip

2b9c36b

support every model

2c1bb6c

Fixed broken tests.

2bb5e62

[2/n] scheduler changes

4d6a05f

[2/n] ip

0831f84

[2/n]ip

f31371f

ip

78bb887

Merge branch 'chunked-prefill-3' into chunked-prefill-scheduler

b9d93c5

[2/n] ip

42dd362

seems to work.

74ac900

Merge branch 'chunked-prefill-3' into chunked-prefill-scheduler

e3afc25

[2/n] ip

6141885

.

71bdada

ip?

d4c3b5d

block tables updated correctly

baef7c6

Merge branch 'chunked-prefill-3' into chunked-prefill-scheduler

d503a22

hopefully tests pass

a12ec68

Merge branch 'chunked-prefill-3' into chunked-prefill-scheduler

85760db

[2/n] update sequence data

e40bc45

[2/n] add prefill range apis

d85670f

Merge branch 'main' into chunked-prefill-3

0d8785f

.

08c8541

ip

3bac9af

add data.

0ca1284

ip

2487bda

simon-mo reviewed Mar 27, 2024

View reviewed changes

benchmarks/benchmark_latency.py Outdated Show resolved Hide resolved

vllm/sequence.py Outdated Show resolved Hide resolved

rkooo567 added the action-required label Mar 27, 2024

rkooo567 added 2 commits March 26, 2024 18:05

remove num chunked prefill from seq group metadata

e0d301c

change apis

5e0f87e

WoosukKwon reviewed Mar 27, 2024

View reviewed changes

cleaned

6e72648

rkooo567 commented Mar 27, 2024

View reviewed changes

rkooo567 added 2 commits March 26, 2024 23:33

now working

4f869be

update with new apis

4f63c57

rkooo567 commented Mar 27, 2024

View reviewed changes

working!

5c3abf4

rkooo567 removed the action-required label Mar 27, 2024

rkooo567 added 2 commits March 27, 2024 06:49

fixed

66f3fcf

Merge branch 'main' into chunked-prefill-scheduler-data-update

9c12d8e

simon-mo approved these changes Mar 28, 2024

View reviewed changes

rkooo567 added the action-required label Mar 28, 2024

Addressed code review.

9d4b65c

rkooo567 removed the action-required label Mar 28, 2024

rkooo567 added 3 commits March 28, 2024 06:24

Merge branch 'main' into chunked-prefill-scheduler-data-update

54a58b2

fix tests.

9bdb9dc

fixed a bug

88126a9

simon-mo merged commit b51c1cc into vllm-project:main Mar 28, 2024
33 checks passed

xjpang pushed a commit to xjpang/vllm that referenced this pull request Mar 31, 2024

[2/N] Chunked prefill data update (vllm-project#3538)

45ebd55

simon-mo mentioned this pull request Apr 4, 2024

[Roadmap] vLLM Roadmap Q2 2024 #3861

Open

64 tasks

AaronFriel mentioned this pull request Apr 6, 2024

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

Open

dtrifiro mentioned this pull request May 15, 2024

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

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[2/N] Chunked prefill data update #3538

[2/N] Chunked prefill data update #3538

rkooo567 commented Mar 21, 2024 •

edited

WoosukKwon Mar 27, 2024

rkooo567 Mar 27, 2024

rkooo567 commented Mar 27, 2024 •

edited

rkooo567 Mar 27, 2024

rkooo567 Mar 27, 2024

rkooo567 left a comment

simon-mo left a comment

simon-mo Mar 28, 2024

rkooo567 Mar 28, 2024

simon-mo Mar 28, 2024

rkooo567 Mar 28, 2024

rkooo567 commented Mar 28, 2024

rkooo567 commented Mar 28, 2024

		@@ -27,18 +27,40 @@ class PreemptionMode(enum.Enum):
		RECOMPUTE = enum.auto()


		class ScheduledSequenceGroup:

[2/N] Chunked prefill data update #3538

[2/N] Chunked prefill data update #3538

Conversation

rkooo567 commented Mar 21, 2024 • edited

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

rkooo567 commented Mar 27, 2024 • edited

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

simon-mo 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 commented Mar 28, 2024

rkooo567 commented Mar 28, 2024

rkooo567 commented Mar 21, 2024 •

edited

rkooo567 commented Mar 27, 2024 •

edited