[Misc] Recompute scheduler upgrade to vLLM 0.18.0

[Angazenn]

What this PR does / why we need it?

The current RecomputeScheduler is aligned to Scheduler in vLLM v0.18.0. Since upstream vLLM has upgraded to v0.18.0, we also need to upgrade RecomputeScheduler to pick up missing updates.

Does this PR introduce any user-facing change?

How was this patch tested?

• vLLM version: v0.18.0
• vLLM main: vllm-project/vllm@35141a7

[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 focuses on upgrading the vllm_ascend recompute scheduler to be compatible with vLLM 0.18.0. The changes involve a comprehensive refactoring of the request scheduling pipeline, including improvements to asynchronous KV cache management, more flexible handling of encoder inputs, and the introduction of scheduler pausing capabilities. These updates aim to enhance the scheduler's efficiency, robustness, and adaptability to advanced KV cache configurations, ensuring seamless integration with the latest vLLM features.

Highlights

• Scheduler Logic Update: The core scheduling logic has been updated to align with the latest vLLM 0.18.0, affecting how requests are managed and processed.
• KV Cache Management Refinement: The _update_waiting_for_remote_kv method was refactored to streamline KV transfer state updates, with readiness checks now handled by the parent scheduler.
• Enhanced KV Cache Configuration: Support for Ascend-specific multi-group KV cache configurations and re-computation of the last token for full prompt hits was added.
• Scheduler Pausing and Request Queue Improvements: The scheduler now respects a PauseState, and the handling of waiting requests has been significantly refactored using step_skipped_waiting and self.skipped_waiting queues for more robust blocked status promotion.
• Encoder Cache Integration: Encoder cache allocation now includes updates via an ec_connector, and the calculation of encoder tokens is more dynamic.
• New Block Zeroing Mechanism: A new mechanism for tracking and zeroing new KV cache blocks has been introduced in the scheduler output.

🧠 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 refactors the RecomputeScheduler to improve handling of waiting requests, particularly for asynchronous KV transfers and paused states. It introduces a new skipped_waiting queue and a more robust mechanism for promoting blocked requests, while also updating the _update_waiting_for_remote_kv method and adding logic for PauseState. The encoder cache allocation logic for encoder-decoder models has been refined by correctly calculating num_encoder_tokens for waiting requests. However, a critical bug was identified where the num_encoder_tokens calculation for cross-attention block allocation is missing for RUNNING requests, which will lead to incorrect KV cache allocation in encoder-decoder models and needs to be addressed.

[gemini-code-assist]

The logic to calculate num_encoder_tokens for cross-attention block allocation has been added for WAITING requests. However, this same logic is missing for RUNNING requests. The allocate_slots call for running requests (around line 327) is not passed num_encoder_tokens. This will result in incorrect KV cache allocation for running requests in encoder-decoder models. A similar logic to calculate and pass num_encoder_tokens should be added for the RUNNING requests loop to fix this bug.

[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.