[Chunked Prefill][4/n] Chunked prefill scheduler. #3853

rkooo567 · 2024-04-04T16:05:53Z

This is a part of the RFC. #3130

This PR

fixes a regression caused from [3/N] Refactor scheduler for chunked prefill scheduling #3550 where RUNNING num_batched tokens are added before running prefill. It is fixed & regression test is added test_scheduler_prefill_prioritized.
Update SchedulingBudget API to include request_ids to dedup add/subtract budgets. It is to make APIs more order agonistic.
From this PR, running and swapped both now can include prefill requests (chunked prefill). To figure out if the seq group is prefill vs decode, the stage is added to SequenceData class.
enable_chunk to support chunking prefills for _schedule_prefills, schedule_running, and _schedule_swapped.
Support chunked prefill scheduling algorithm. The algorithm is as follow;

1. Schedule all decodes
2. Schedule all chunked prefills (running prefills)
3. Schedule swapped
4. Schedule new prefills

Since chunked prefill requests are now a part of running, preemption & swapping works with chunked prefill. There are test cases covering it in a new test file test_chunked_prefill_scheduler.py.

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!

vllm/core/policy.py

simon-mo

does this mean we can do end to end test now?

vllm/core/policy.py

vllm/sequence.py

simon-mo · 2024-04-04T21:59:04Z

vllm/engine/llm_engine.py


        for scheduled_seq_group, outputs in zip(scheduled_seq_groups, output):
            seq_group = scheduled_seq_group.seq_group
-            token_chunk_size = scheduled_seq_group.token_chunk_size


why removed?

I made it handled in scheduler.py. Basically as soon as we confirm it is scheduled, we update num computed tokens (which is consistent to moving waiting item to running within a scheduler).

vllm/core/scheduler.py

simon-mo · 2024-04-04T22:06:37Z

vllm/core/scheduler.py

+                            seq_group=seq_group,
+                            token_chunk_size=num_running_tokens))
+                else:
+                    assert num_running_tokens == 1, num_running_tokens


would this break spec decode because the rejection sampling takes multiple token @cadedaniel

I think when the feature is working e2e, we can remove the assert?

rkooo567 · 2024-04-04T23:08:14Z

does this mean we can do end to end test now?

We will need attention level change which will come in the next PR. That'd be the last one

rkooo567 · 2024-04-04T16:08:36Z

vllm/core/scheduler.py

+    def remaining_token_budget(self):
+        return self.token_budget - self.num_batched_tokens
+
+    def add_num_batched_tokens(self, seq_group: SequenceGroup,


It's required so that we can avoid having duplicated budget add/subtract

rkooo567 · 2024-04-04T16:11:31Z

vllm/core/scheduler.py

        )
+        # Make sure we include num running seqs before scheduling prefill,
+        # so that we don't schedule beyond max_num_seqs for prefill.
+        for seq_group in self.running:


this was a bug (we should not include num_batched tokens in the beginning). it is fixed & regression tests are added.

vllm/core/policy.py

vllm/sequence.py

rkooo567 · 2024-04-04T23:14:26Z

vllm/engine/llm_engine.py


        for scheduled_seq_group, outputs in zip(scheduled_seq_groups, output):
            seq_group = scheduled_seq_group.seq_group
-            token_chunk_size = scheduled_seq_group.token_chunk_size


I made it handled in scheduler.py. Basically as soon as we confirm it is scheduled, we update num computed tokens (which is consistent to moving waiting item to running within a scheduler).

simon-mo · 2024-04-04T23:28:06Z

vllm/sequence.py

+        """
+        num_uncomputed_tokens = self.data.get_num_uncomputed_tokens()
+        if self.data.stage == SequenceStage.DECODE:
+            assert num_uncomputed_tokens == 1


where is this set? get_num_uncomputed_tokens in RequestMetrics only returns return self.get_len() - self.get_num_computed_tokens()

Hmm yeah, this assumes that a new output token is added to sequence (so that get_len is incremented by 1) before update happens. I just removed the assert here

vllm/core/scheduler.py

simon-mo

Approving. The PR is mostly additive and some slight refactoring over existing code path to enable the chunked prefill path.

vllm/core/scheduler.py

AgrawalAmey · 2024-04-04T21:13:19Z

vllm/core/scheduler.py

+        self._requeset_ids_num_curr_seqs.add(req_id)
+        self._num_curr_seqs += num_curr_seqs
+
+    def subtract_num_seqs(self, seq_group: SequenceGroup, num_curr_seqs: int):


i think the two add/substract methods should be combined

unfortunately, for existing logic, we only add budget for num_seqs. So we need a separate method here...

AgrawalAmey · 2024-04-04T21:13:37Z

vllm/core/scheduler.py

+    def add_num_batched_tokens(self, seq_group: SequenceGroup,
+                               num_batched_tokens: int):
+        req_id = seq_group.request_id
+        if req_id in self._requeset_ids_num_batched_tokens:


when would this happen? I think this should raise an exception

same case as #3853 (comment)

vllm/core/scheduler.py

AgrawalAmey · 2024-04-04T21:18:29Z

vllm/core/scheduler.py

+                seq_group, SequenceStatus.RUNNING, enable_chunking, budget)
+            # We can have up to 1 running prefill at any given time in running
+            # queue, which means we can guarantee chunk size is at least 1.
+            assert num_running_tokens != 0


wha if budget is zero?

This can't happen now. (because for normal case, if prefill is scheduled, this is not called. In chunked prefill, decoding is called first. I don't think we need to handle num_max_batched_tokens=0). So I will leave it in a current way and not handle this case (we can remove assert if there's other cases that this can happen)

rkooo567 added 4 commits April 3, 2024 23:21

code complete. wip test

6b71130

ip

6ab059c

ip

2ae3444

Done

07c35b9

rkooo567 mentioned this pull request Apr 4, 2024

[RFC] Upstream Chunked Prefill #3130

Closed

done

29b1a26

AgrawalAmey reviewed Apr 4, 2024

View reviewed changes

vllm/core/policy.py Outdated Show resolved Hide resolved

simon-mo requested review from simon-mo and zhuohan123 April 4, 2024 21:53

simon-mo reviewed Apr 4, 2024

View reviewed changes

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

bcbad91

rkooo567 commented Apr 4, 2024

View reviewed changes

Fixed

81e8dc4

simon-mo reviewed Apr 4, 2024

View reviewed changes

rkooo567 added 3 commits April 4, 2024 16:39

addressed 2

7ee4655

removed assert

afa0df2

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

24fe6dc

simon-mo approved these changes Apr 5, 2024

View reviewed changes

rkooo567 added 5 commits April 4, 2024 18:57

test fix

df86440

revert basic corr test

1941fec

done

4f64529

Fixed an issue

b570f96

fixed a bug

f93e677

AgrawalAmey reviewed Apr 5, 2024

View reviewed changes

rkooo567 added 5 commits April 4, 2024 23:22

done

150beab

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

020b31f

fix

e7c0d6c

trial

8ba003e

test add

26193dd

fixed tests

705313e

simon-mo merged commit 18de883 into vllm-project:main Apr 5, 2024

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

[Chunked Prefill][4/n] Chunked prefill scheduler. (vllm-project#3853)

408bfbd

dtrifiro mentioned this pull request May 15, 2024

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

Merged

toslunar mentioned this pull request Jun 18, 2024

[Bugfix] [Core] don't schedule prefill if freeing kv cache #5633

Closed

toslunar mentioned this pull request Feb 19, 2025

fix(chunked prefill): don't schedule prefill if freeing kv cache #13539

Closed

Uh oh!

[Chunked Prefill][4/n] Chunked prefill scheduler. #3853

[Chunked Prefill][4/n] Chunked prefill scheduler. #3853

Uh oh!

Conversation

rkooo567 commented Apr 4, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Title and Classification

Code Quality

Notes for Large Changes

What to Expect for the Reviews

Thank You

Uh oh!

Uh oh!

simon-mo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rkooo567 commented Apr 4, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

simon-mo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rkooo567 commented Apr 4, 2024 •

edited

Loading