[V1] Support long_prefill_token_threshold in v1 scheduler

[houseroad]

To address #14003

To support concurrent partial prefill, we need to define long_prefill_token_threshold. Then the tokens from different prefill requests can be added. Since there is no difference between prefill and decode, just need to control the token threshold.

pytest tests/v1/core/test_scheduler.py
e2e test: VLLM_USE_V1=1 pytest tests/v1/core/test_scheduler_e2e.py

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[plops655]

I am confused as to whether this PR properly enables chunked prefill for v1. In the original issue, my understanding is that concurrent partial prefills are used only when chunked prefills are enabled. Also, within the original scheduler code, we made sure that prefills that are scheduled to finish so that these sequences are prioritized over unfinished prefills upon transitioning to decode stage.

# Because multiple prefills may be running concurrently, we need to
        # make sure that prefills which are scheduled to finish are listed
        # before those that won't. This is so that on the next scheduling
        # iteration when they have transitioned to the decode stage, they are
        # properly prioritized over sequences that are still in the prefill
        # stage.
        self.running.extend(
            self._order_finishing_prefills_first(
                running_scheduled.prefill_seq_groups))
        self.running.extend([s.seq_group for s in prefills.seq_groups])

in vllm/core/scheduler.py _schedule_chunked_prefill

[houseroad]

This PR is only for concurrent partial prefills, not for chunked prefill. The assumption is that v1 always has the chunked prefill, so no difference between prefill and decoding.

Feel free to suggest any fix if something is wrong. :-)

[plops655]

Thanks for responding. I have one more question I would be grateful to know, as I am relatively new to vLLM. In v0, chunked prefill is explicitly run in _schedule_chunked_prefill. This function creates a PartiallPrefillMetadata object for controlling the number of concurrent partial prefills, and batches both chunked prefills and decodes together.

I do not see any explicit chunking happening in the v1 schedule() code, and was wondering if you could enlighten me as to how chunked prefill is the default in v1.

[comaniac]

Thanks for responding. I have one more question I would be grateful to know, as I am relatively new to vLLM. In v0, chunked prefill is explicitly run in _schedule_chunked_prefill. This function creates a PartiallPrefillMetadata object for controlling the number of concurrent partial prefills, and batches both chunked prefills and decodes together.

I do not see any explicit chunking happening in the v1 schedule() code, and was wondering if you could enlighten me as to how chunked prefill is the default in v1.

Chunked prefill is the first class citizen in v1 scheduler so there's no explicit flag to configure it. There are some places you can observe this behavior:

• When determining the number of new tokens to schedule (https://github.com/vllm-project/vllm/blob/main/vllm/v1/core/sched/scheduler.py#L290), we don't enforce the entire prompt to be scheduled at once. This is because of "chunked" prefill.
• When scheduling running requests, we can schedule arbitrary number of tokens, and don't check whether it's prompt tokens or decode tokens. This is also because of the support of chunked prefill.

[comaniac]

LGTM!
Also cc @joerunde @njhill

[njhill]

Just looking at this now.. small nit that this could be simplified slightly

            if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
                num_new_tokens = self.scheduler_config.long_prefill_token_threshold

[comaniac]

Good point. Fixed in #15307