Refactor / Unify event loop across PD-Disagg, Overlap, DP-Attn cases

[hnyls2002]

This PR replaces #9618.

Modifications

• Introduce the prebuilt_extend forward mode, which is used in the decode event loop in pd-disaggregation mode.
• Unify the prepare_mlp_sync logic, place this function inside get_next_disagg_decode_batch_to_run.

The all-gather logic for decode batch, empty batch, and prebuilt batch is as follows

• Decode (dp0) + Empty (dp1) -> Docode (dp0) + Idle (dp1)
• Decode (dp0) + Prebuilt (dp1) -> Decode (dp0) + Idle (dp1, but returns prebuilt batch)
• Prebuilt (dp0) + Empty (dp1) -> No forward (but prebuilt batch returns)
• Decode (dp0) + Empty (dp1) + Prebuilt (dp2) -> Decode (dp0) + Idle (dp1) + Idle (dp2, but returns prebuilt)

[gemini-code-assist]

Summary of Changes

Hello @hnyls2002, 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 enhancing the maintainability and efficiency of the SGLang runtime's scheduling mechanisms. By unifying event loop patterns and centralizing key operational checks, the changes aim to create a more robust and easier-to-understand system for managing different batch processing scenarios, particularly in disaggregated and parallel environments. The introduction of a specialized forward mode for decode operations further refines how the system handles KV cache readiness.

Highlights

• Event Loop Unification: The event loop logic across different disaggregation and attention cases (PD-Disagg, Overlap, DP-Attn) has been refactored and unified, reducing code duplication and simplifying the overall flow.
• Centralized MLP Sync Handling: The require_mlp_sync check is now handled centrally by storing its state as self.require_mlp_sync within the Scheduler instance, rather than being repeatedly evaluated in various event loops.
• New Forward Mode for Disaggregated Decode: A new ForwardMode.FAKE_EXTEND has been introduced to specifically manage batches where KV cache is ready for decoding in disaggregated decode workers, streamlining their processing.
• Refactored Batch Execution: A new method, _run_batch_fake_extend, has been added to the Scheduler to encapsulate the logic for handling FAKE_EXTEND batches, leading to cleaner and more modular batch processing.
• Simplified Decode Event Loops: The event_loop_normal_disagg_decode and event_loop_overlap_disagg_decode methods have been simplified by removing specific is_extend() checks and prepare_mlp_sync_flag logic, as these are now handled by the new unified mechanisms.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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 event loops for prefill-decode disaggregation, overlap scheduling, and data-parallel attention to unify their logic. This is a significant and positive change that improves code clarity and maintainability. The introduction of a new FAKE_EXTEND forward mode and centralizing the handling of different batch types in Scheduler.run_batch is a good approach. The cleanup of require_mlp_sync calls is also a welcome improvement. I have found one potential issue regarding missing output streaming for the new FAKE_EXTEND mode, which could lead to the first token being dropped in disaggregated mode.

[hnyls2002]

All comments are resolved in #12948