[Patch] Support balance scheduling

[GDzhu01]

What this PR does / why we need it?

This PR introduces a "balance scheduling" feature, enabled by the VLLM_ASCEND_BALANCE_SCHEDULING environment variable. This feature adjusts the scheduling logic to better balance the load across data-parallel workers, preventing a single worker from blocking scheduling for others. This can improve overall throughput.

Additionally, this PR includes a number of other updates and fixes to the scheduler, syncing it with a more recent version of the upstream vLLM scheduler. These changes include:

• Handling for paused scheduler state.
• Support for Mamba block-aligned splits.
• Handling for streaming requests.
• Refinements in preemption logic and resource management (KV cache, encoder cache).
• General code refactoring for clarity and correctness.

Fixes #

Does this PR introduce any user-facing change?

Yes, this PR introduces a new feature controlled by the VLLM_ASCEND_BALANCE_SCHEDULING environment variable. When enabled, the scheduling behavior changes, which could affect performance and request throughput.

How was this patch tested?

CI passed. Further testing should be done to validate the performance and correctness of the new scheduling logic under various workloads, with and without the feature flag enabled.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for balance scheduling within the VLLM Ascend platform. The changes enable more granular control over the scheduler's operation, particularly by allowing it to be paused and by refining how requests are processed based on their state and specific model requirements like Mamba block alignment. This aims to improve resource management and scheduling efficiency, especially in complex distributed or streaming scenarios.

Highlights

• Conditional Balance Scheduling: Introduced a mechanism to conditionally enable balance scheduling functionality based on the VLLM_ASCEND_BALANCE_SCHEDULING environment variable, ensuring the feature is only active when explicitly desired.
• Scheduler Pause State Integration: Integrated a PauseState into the scheduler, allowing the system to halt token scheduling when in a PAUSED_ALL state, providing better control over resource allocation.
• Refined Request Handling and KV Cache Management: Improved the handling of requests in WAITING_FOR_REMOTE_KVS and WAITING_FOR_STREAMING_REQ states, and updated KV cache allocation logic to include Mamba-specific block alignment and proper encoder cache freeing upon allocation failure.
• Enhanced Prefix Cache Statistics: Added detailed tracking for connector prefix cache queries and hits, providing better insights into cache performance and preemption scenarios.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for balanced scheduling, controlled by the VLLM_ASCEND_BALANCE_SCHEDULING environment variable. The changes modify the scheduler logic to improve load balancing across data-parallel workers. The PR also includes several other updates and refactorings to the scheduler, aligning it with more recent upstream changes, including support for Mamba models and streaming requests.

The pull request title and description could be improved to better reflect the changes and adhere to the repository's contribution guidelines. Here is a suggestion:

Suggested PR Title:

[Scheduler][Feature] Support balance scheduling and sync scheduler with upstream

Suggested PR Summary:

### What this PR does / why we need it?
This PR introduces a "balance scheduling" feature, enabled by the `VLLM_ASCEND_BALANCE_SCHEDULING` environment variable. This feature adjusts the scheduling logic to better balance the load across data-parallel workers, preventing a single worker from blocking scheduling for others. This can improve overall throughput.

Additionally, this PR includes a number of other updates and fixes to the scheduler, syncing it with a more recent version of the upstream vLLM scheduler. These changes include:
- Handling for paused scheduler state.
- Support for Mamba block-aligned splits.
- Handling for streaming requests.
- Refinements in preemption logic and resource management (KV cache, encoder cache).
- General code refactoring for clarity and correctness.

Fixes #

### Does this PR introduce _any_ user-facing change?
Yes, this PR introduces a new feature controlled by the `VLLM_ASCEND_BALANCE_SCHEDULING` environment variable. When enabled, the scheduling behavior changes, which could affect performance and request throughput.

### How was this patch tested?
CI passed. Further testing should be done to validate the performance and correctness of the new scheduling logic under various workloads, with and without the feature flag enabled.