Sync cuda graph runners

[merrymercy]

No description provided.

[gemini-code-assist]

Hello @merrymercy, 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-code-assist here to provide a summary of this pull request. Based on the title "Sync cuda graph runners" and the changes, it appears this PR aims to synchronize and make minor cleanups to the CUDA graph runner implementations, specifically for the 'eagle_draft' and 'eagle_draft_extend' runners. The main change involves ensuring that a tensor sum is converted to a scalar value using .item() where needed, and there's a small code style improvement.

Highlights

• Data Type Consistency: Ensure seq_lens_sum is a scalar by calling .item() on the tensor sum result in both eagle_draft_cuda_graph_runner.py and eagle_draft_extend_cuda_graph_runner.py. This likely prevents potential issues with tensor operations expecting scalar inputs.
• Code Style Cleanup: Remove redundant parentheses from an if condition in eagle_draft_cuda_graph_runner.py for minor code style improvement.

Changelog

• python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
  • Convert seq_lens.sum() to seq_lens.sum().item() on line 130 to get a scalar value.
  • Remove redundant parentheses from the if condition on line 212.
• python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
  • Convert seq_lens.sum() to seq_lens.sum().item() on line 141 to get a scalar value.

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.

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 small but important changes to the CUDA graph runners, primarily by ensuring seq_lens_sum is a Python scalar using .item(). This is excellent practice, especially in the context of CUDA graphs and custom C++/CUDA backends, as it enhances robustness and prevents potential type-related errors. The changes contribute to better code quality and clarity.

Well done on these improvements!

Summary of Findings

• Stylistic Improvement in Conditional Statement: In python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py (line 212), parentheses around conditions in an if statement were removed (e.g., if (A) and (B): became if A and B:). This is a minor stylistic improvement that enhances readability by reducing visual clutter, as Python's operator precedence makes the parentheses unnecessary here. This was a low-severity finding and thus not commented on directly in the code per review settings.

Merge Readiness

The changes in this pull request are beneficial and improve the robustness of the CUDA graph runners. The code appears to be in good shape and ready for merging. As always, ensure all CI checks pass. I am an AI assistant and cannot approve pull requests, so please ensure appropriate human review and approval before merging.