[Core] Chunked Prefill support for Multi Step Scheduling #7814

varun-sundar-rabindranath · 2024-08-23T08:17:03Z

Add Chunked Prefill support for Multi Step Scheduling.

Based on, #7528, this PR adds both the Force Single Step and Ignore Prefill policies.

The Ignore Prefill policy implemented as the default policy with a VLLM_MULTI_STEP_CHUNKED_PREFILL_SINGLE_STEP_POLICY to switch to the Force Single Step policy.

Benchmark :

Please find A100 benchmark results here

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!

github-actions · 2024-08-23T08:17:16Z

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

Comment /ready on the PR
Add ready label to the PR
Enable auto-merge.

🚀

SolitaryThinker · 2024-08-23T08:55:21Z

vllm/core/scheduler.py

@@ -1202,7 +1217,8 @@ def _append_slots(
                the new source and destination block indices for the appended
                slots.
        """
-        num_lookahead_slots = self._get_num_lookahead_slots(is_prefill=False)
+        num_lookahead_slots = self._get_num_lookahead_slots(
+            is_prefill=seq_group.is_prefill())


Spent some time looking into PP. The assertion issue may be related to this change. Something is causing the batch to generate logits for the prefills in the batch and since len(sample_indices) == 0) as they don't perform sampling, the assertion fails on assert logits_applied == logits.shape[0] in _apply_min_tokens_penalty since logits_applied is the sum of all sampled_indices in the batch

I see. Thanks @SolitaryThinker. I believe it should be

num_lookahead_slots = self._get_num_lookahead_slots(is_prefill=seq_group.is_prefill() and seq_group.do_sample)

instead of just

num_lookahead_slots = self._get_num_lookahead_slots(is_prefill=False)

I haven't been able to reproduce this assertion yet. I'll keep at it. When you find some time, can you try this out as well ? Thanks.

976d032 fixes the PP issue. The SamplingMetadata objects were being clobbered due to the SamplingMetadataCache reset.

comaniac

Thanks for the great work! Leave some comments.

comaniac · 2024-08-23T20:15:11Z

vllm/engine/async_llm_engine.py

+                if seq_group.state.remaining_steps > 0:
+                    seq_group.finish_step()


Could you comment on this?

comaniac · 2024-08-23T20:16:24Z

vllm/engine/async_llm_engine.py

+            # The prefill sequences's remaining_step is 1 when they are
+            # scheduled initially. After the first step their remaining_step
+            # becomes 0.
+            if any([sgml.state.remaining_steps not in [0, 1, remaining_steps] \


Suggested change

if any([sgml.state.remaining_steps not in [0, 1, remaining_steps] \

if any([sgml.state.remaining_steps not in (0, 1, remaining_steps) \

comaniac · 2024-08-23T20:20:00Z

vllm/envs.py

@@ -60,6 +60,7 @@
    VLLM_ALLOW_ENGINE_USE_RAY: bool = False
    VLLM_PLUGINS: Optional[List[str]] = None
    VLLM_TORCH_PROFILER_DIR: Optional[str] = None
+    VLLM_MULTI_STEP_CHUNKED_PREFILL_SINGLE_STEP_POLICY: bool = False


Would that be better if we make it more extensible like the following?

Suggested change

VLLM_MULTI_STEP_CHUNKED_PREFILL_SINGLE_STEP_POLICY: bool = False

VLLM_MULTI_STEP_CHUNKED_PREFILL_POLICY: str = "let-prefill-wait"

And another policy would be "single-step-with-prefill"

comaniac · 2024-08-23T20:22:35Z

tests/multi_step/test_correctness.py

+class MultiStepChunkedPrefillPolicy(Enum):
+    # When prompt and decode sequences are scheduled together,
+    # the DEFAULT policy is to run the prompt and decodes sequences
+    # together only for the first step and run just the decode sequences
+    # in the rest of the steps.
+    DEFAULT = 1
+    # In FORCE_SINGLE_STEP policy, we force the scheduled sequences to
+    # run a single step and then re-schedule.
+    FORCE_SINGLE_STEP = 2
+    INVALID = 3


We should define this in vLLM instead of tests.

comaniac · 2024-08-23T20:25:22Z

vllm/worker/multi_step_model_runner.py

@@ -208,11 +277,42 @@ def prepare_model_input(
        frozen_model_input = self._base_model_runner.prepare_model_input(
            seq_group_metadata_list, virtual_engine, finished_requests_ids)

+        num_prompts = len(


Change all "prompts" to "prefills".

Suggested change

num_prompts = len(

num_prefills = len(

comaniac · 2024-08-23T20:30:55Z

vllm/worker/model_runner.py

+    def without_prefills(m: "ModelInputForGPUWithSamplingMetadata",
+                        sampling_metadata_decodes: SamplingMetadata) \
+                            -> "ModelInputForGPUWithSamplingMetadata":


Add docstring

"""Remove prefill sequences from this model input metadata. This is used by multi-step runner to execute the rest steps of decode sequences. """

comaniac · 2024-08-23T20:31:32Z

vllm/worker/model_runner.py

+        assert (m.input_tokens is not None)
+        assert (m.input_positions is not None)


Suggested change

assert (m.input_tokens is not None)

assert (m.input_positions is not None)

assert m.input_tokens is not None

assert m.input_positions is not None

comaniac · 2024-08-23T20:33:42Z

tests/multi_step/test_correctness.py

@@ -23,17 +48,36 @@
 ]


-async def completions_with_server_args(prompts: List[str], model_name: str,
-                                       server_cli_args: List[str]):
+class EnvContextManager():


Can you look around the tests and see if we already have a similar utility? If so please reuse it; otherwise please move this to the right place so that other tests in the future can use it.

comaniac · 2024-08-23T20:36:26Z

tests/multi_step/test_correctness.py

+        os.environ.update(self.os_env)
+
+
+async def completions_with_server_args(prompts: List[str],


#7651 also introduces this utility. Please coordinate with @afeldman-nm to better organize these.

@comaniac do you mean #7652?

Oops yes sorry for the typo.

comaniac · 2024-08-23T20:37:26Z

vllm/attention/backends/flash_attn.py

+    # TODO (varun) : Try using decode_metadata here. We hit some asserts in
+    # advance_step - but that seems resolvable.


Can you elaborate more on this? Specifically how it can be solved and what's the plan?

varun-sundar-rabindranath · 2024-08-26T15:37:26Z

tests/multi_step/test_correctness.py

+                                       server_cli_args: List[str],
+                                       with_env: dict = {}):  # noqa: B006
+    # env setup
+    os.environ.update(with_env)


TODO : fix stray update.

varun-sundar-rabindranath · 2024-09-11T18:52:47Z

Bad performance - just multi-step is better than multi-step + chunked-prefill closing this in favor of #8378

varun-sundar-rabindranath · 2024-09-11T18:53:19Z

@tlrmchlsmth @comaniac @SolitaryThinker I'll handle the relevant comments in the new PR

Varun Sundar Rabindranath added 8 commits August 22, 2024 21:17

multi-step + chunked prefill policy 2

8dc48c4

nits

b15ef3a

fix has_remaining_steps

af07307

add comments

d2c1f9c

add force single-step logic and add policy switcher

2b8a6a1

update tests

e137fa2

format

2a44012

reset tests

93cc29f

varun-sundar-rabindranath mentioned this pull request Aug 23, 2024

Varun/multi step chunked prefill #7563

Closed

fix bad asserts

db91e3e

SolitaryThinker reviewed Aug 23, 2024

View reviewed changes

Varun Sundar Rabindranath added 3 commits August 23, 2024 15:28

broadcast sampling_metadata_decodes

047a0ea

format

14bb0e9

fix bad assert

6eac258

comaniac reviewed Aug 23, 2024

View reviewed changes

comaniac mentioned this pull request Aug 23, 2024

[Tracking issue] [Help wanted]: Multi-step scheduling follow-ups #7528

Open

17 tasks

Varun Sundar Rabindranath added 2 commits August 25, 2024 05:43

fix PP : avoid using sampling metadata cache

976d032

format

be7fe61

varun-sundar-rabindranath commented Aug 26, 2024

View reviewed changes

remove stray update

5b29a05

varun-sundar-rabindranath closed this Sep 11, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core] Chunked Prefill support for Multi Step Scheduling #7814

[Core] Chunked Prefill support for Multi Step Scheduling #7814

varun-sundar-rabindranath commented Aug 23, 2024 •

edited

Loading

github-actions bot commented Aug 23, 2024

SolitaryThinker Aug 23, 2024 •

edited

Loading

varun-sundar-rabindranath Aug 23, 2024

varun-sundar-rabindranath Aug 26, 2024

comaniac left a comment

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

comaniac Aug 23, 2024

tlrmchlsmth Aug 26, 2024

comaniac Aug 26, 2024

comaniac Aug 23, 2024

varun-sundar-rabindranath Aug 26, 2024

varun-sundar-rabindranath commented Sep 11, 2024

varun-sundar-rabindranath commented Sep 11, 2024

		if seq_group.state.remaining_steps > 0:
		seq_group.finish_step()

	if any([sgml.state.remaining_steps not in [0, 1, remaining_steps] \
	if any([sgml.state.remaining_steps not in (0, 1, remaining_steps) \

	VLLM_MULTI_STEP_CHUNKED_PREFILL_SINGLE_STEP_POLICY: bool = False
	VLLM_MULTI_STEP_CHUNKED_PREFILL_POLICY: str = "let-prefill-wait"

		assert (m.input_tokens is not None)
		assert (m.input_positions is not None)

		os.environ.update(self.os_env)


		async def completions_with_server_args(prompts: List[str],

		# TODO (varun) : Try using decode_metadata here. We hit some asserts in
		# advance_step - but that seems resolvable.

[Core] Chunked Prefill support for Multi Step Scheduling #7814

[Core] Chunked Prefill support for Multi Step Scheduling #7814

Conversation

varun-sundar-rabindranath commented Aug 23, 2024 • edited Loading

Please find A100 benchmark results here

PR Title and Classification

Code Quality

Notes for Large Changes

What to Expect for the Reviews

Thank You

github-actions bot commented Aug 23, 2024

SolitaryThinker Aug 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

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

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

varun-sundar-rabindranath commented Sep 11, 2024

varun-sundar-rabindranath commented Sep 11, 2024

varun-sundar-rabindranath commented Aug 23, 2024 •

edited

Loading

SolitaryThinker Aug 23, 2024 •

edited

Loading