Minor refactor two-batch overlap

[fzyzcjy]

Motivation

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @fzyzcjy, 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!

Summary of Changes

Hello team, Gemini here to provide a summary of this pull request. This PR focuses on a minor refactor related to the 'two-batch overlap' (TBO) feature, specifically concerning its interaction with CUDA graphs. The primary goal is to consolidate and organize the logic responsible for preparing and replaying CUDA graphs when TBO is enabled, moving it into a dedicated plugin class.

Highlights

• Refactor Two-Batch Overlap (TBO) CUDA Graph Logic: The main change is the refactoring of code related to how the two-batch overlap feature interacts with CUDA graph capture and replay.
• Introduce TboCudaGraphRunnerPlugin: A new class, TboCudaGraphRunnerPlugin, is introduced in two_batch_overlap.py to encapsulate the TBO-specific logic for the CudaGraphRunner.
• Consolidate Split Index Computation: The logic for computing TBO split sequence and token indices specifically for CUDA graph replay is moved into a new helper function compute_split_indices_for_cuda_graph_replay in two_batch_overlap.py.
• Delegate CUDA Graph Preparation to Plugin: The CudaGraphRunner now instantiates TboCudaGraphRunnerPlugin and delegates the TBO-specific steps during CUDA graph capture and replay preparation to methods on this plugin.

Changelog

Click here to see the changelog

• python/sglang/srt/layers/attention/tbo_backend.py
  • Removed manual computation of tbo_split_seq_index and tbo_split_token_index (lines 128-140).
  • Replaced manual computation with a call to the new consolidated function two_batch_overlap.compute_split_indices_for_cuda_graph_replay (lines 126-131).
  • Removed unused import of ForwardMode (line 122).
• python/sglang/srt/model_executor/cuda_graph_runner.py
  • Renamed imported utility class from TboCudaGraphRunnerUtils to TboCudaGraphRunnerPlugin (line 43).
  • Added an instance of the new TboCudaGraphRunnerPlugin as self.tbo_plugin in the __init__ method (line 259).
  • Removed direct call to TboCudaGraphRunnerUtils.compute_tbo_split_seq_index (lines 484-486).
  • Removed direct call to TboForwardBatchPreparer.prepare (line 489).
  • Delegated TBO capture logic to self.tbo_plugin.capture_one_batch_size (line 487).
  • Modified replay_prepare to delegate TBO preparation logic to self.tbo_plugin.replay_prepare (lines 584-588).
• python/sglang/srt/two_batch_overlap.py
  • Added new function compute_split_indices_for_cuda_graph_replay to consolidate logic for CUDA graph split index computation (lines 88-105).
  • Renamed class TboCudaGraphRunnerUtils to TboCudaGraphRunnerPlugin (line 111).
  • Moved TBO capture logic (computing split index and calling TboForwardBatchPreparer.prepare) into the new TboCudaGraphRunnerPlugin.capture_one_batch_size method (lines 115-127).
  • Added a placeholder TboCudaGraphRunnerPlugin.replay_prepare method (lines 129-135).

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

---

Code moves around,
Logic finds new ground,
A plugin takes hold,
A story unfolds,
Refactorings abound.

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 effectively refactors the two-batch overlap (TBO) logic by introducing a TboCudaGraphRunnerPlugin and centralizing related computations. This improves modularity and makes the TBO interactions within the CUDA graph runner and attention backend clearer.

The core functionality appears to be preserved. The main points for discussion are the TODO comments in the new plugin, which might indicate areas for further development or clarification.

Summary of Findings

• Incomplete TboCudaGraphRunnerPlugin methods: The __init__ and replay_prepare methods in the new TboCudaGraphRunnerPlugin class contain TODO comments with pass statements. It would be good to clarify the intended logic or future plans for these sections to ensure the plugin's role is fully realized.
• Correctness of TBO logic movement: Verify that moving the TBO split index computation and batch preparation logic to the TboCudaGraphRunnerPlugin does not introduce any regressions or unexpected behavior, especially in the capture_one_batch_size and replay_prepare methods.

Merge Readiness

The refactoring is a positive step towards better code organization for the TBO feature. However, before merging, it would be beneficial to address the TODO comments in TboCudaGraphRunnerPlugin by either implementing the planned logic or clarifying if these are placeholders for future work beyond the scope of this PR. Additionally, confirm the correctness of the TBO logic movement and parameter passing in the capture_one_batch_size and replay_prepare methods. Depending on the clarification, changes might be needed. As a reviewer, I am not authorized to approve pull requests, so please ensure further review and approval from maintainers.

[gemini-code-assist]

The __init__ method for the new TboCudaGraphRunnerPlugin currently contains a pass statement with a TODO add logic here comment.

Could you clarify if there's any initialization logic planned for this plugin in the near future, or if this TODO is more of a placeholder for potential future state management? If no initialization is currently needed, this is fine, but understanding the intent would be helpful.

[gemini-code-assist]

The replay_prepare method in TboCudaGraphRunnerPlugin is newly introduced as a hook in CudaGraphRunner.replay_prepare. It currently contains a pass statement with a TODO add logic here comment after checking enable_two_batch_overlap.

What is the intended purpose of this method? Is there specific TBO-related preparation logic that needs to be executed during CUDA graph replay that is planned to be added here? If so, this PR might be setting up the structure for it. Clarifying the scope of this TODO would be beneficial.

[gemini-code-assist]

The original code computes tbo_split_seq_index here. Now that this logic is moved to TboCudaGraphRunnerPlugin, is it correct to remove it from here? Please confirm that this change aligns with the intended behavior.

[gemini-code-assist]

This section introduces a call to self.tbo_plugin.replay_prepare. It's important to ensure that the parameters passed to this function (forward_mode, bs, num_token_non_padded) are the correct and expected values for the TBO logic during replay preparation. Double-check these parameters to prevent any unexpected behavior.