Sync the changes on 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 titled "Sync up the changes on cuda graph runners". Based on the title and the code changes, this PR appears to synchronize the CUDA graph runner implementation with recent updates or requirements, particularly concerning attention backends, speculative decoding (Eagle/Eagle3), data parallelism, and error handling.

The changes primarily involve refining the initialization logic within the CudaGraphRunner class, adjusting how certain buffers and speculative decoding parameters are handled, improving error messaging for capture failures, and making a fix related to pipeline parallelism replay.

Highlights

• Attention Backend Initialization: Simplified the initialization of the attention backend's CUDA graph state by removing the conditional logic for the flashmla backend. The initialization now consistently uses self.max_num_token (line 216).
• Speculative Decoding Refinements: Removed the explicit initialization of the self.hidden_states tensor for Eagle/Eagle3 speculative algorithms in the __init__ method (lines 259-260 removed) and also removed the copying of hidden states during replay_prepare (line 593-594 removed). This suggests a change in how hidden states are managed for speculation within the CUDA graph context. Additionally, the default custom_mask for Eagle verification input was changed from torch.zeros to torch.ones (line 629), and a topk parameter was added to the EagleDraftInput (line 640).
• Gathered Buffer Initialization: Modified the condition for initializing the self.gathered_buffer. It is now only initialized if self.enable_dp_attention is true, removing the dependency on self.enable_sp_layernorm (line 270).
• Improved Error Messaging: Replaced the hardcoded multi-line error message for CUDA graph capture failures with a reference to a new constant CUDA_GRAPH_CAPTURE_FAILED_MSG (line 288), which is defined later in the file (lines 648-655).
• Pipeline Parallelism Replay Fix: Ensured that pp_proxy_tensors are cloned before being passed to the forward function during CUDA graph replay in run_once (lines 477-479). This prevents potential unintended modifications to the original tensors used for replay setup.

Changelog

Click here to see the changelog

• python/sglang/srt/model_executor/cuda_graph_runner.py
  • Simplified attention backend initialization logic in __init__ (DIFF 1).
  • Removed explicit self.hidden_states initialization for Eagle/Eagle3 in __init__ (DIFF 2).
  • Adjusted condition for initializing self.gathered_buffer in __init__ (DIFF 3).
  • Replaced hardcoded CUDA graph capture error message with CUDA_GRAPH_CAPTURE_FAILED_MSG constant (DIFF 4).
  • Added a blank line (DIFF 5).
  • Reordered arguments in the call to tbo_plugin.capture_one_batch_size (DIFF 6).
  • Cloned pp_proxy_tensors before passing to forward in run_once (DIFF 7).
  • Removed copying of hidden states during replay_prepare (DIFF 8).
  • Changed default custom_mask from zeros to ones for Eagle verification in get_spec_info (DIFF 9).
  • Added topk parameter to EagleDraftInput and defined CUDA_GRAPH_CAPTURE_FAILED_MSG constant (DIFF 10).

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 several changes to the CUDA graph runner, primarily focused on speculative decoding (Eagle) and attention backend interactions. Key modifications include adjustments to CUDA graph state initialization, handling of hidden states for speculative inference, and updates to the EagleVerifyInput parameters. While some changes appear to be cleanups or refactorings, a few raise concerns regarding correctness and require further clarification.

Summary of Findings

• CUDA Graph State Initialization for FlashMLA: The init_cuda_graph_state for the attention backend is now unconditionally called with self.max_num_token. This appears problematic for the flashmla backend, which seems to expect its max_bs parameter to be the actual maximum batch size, not the total number of tokens. This could lead to incorrect memory allocation or behavior.
• Handling of Hidden States for Speculative Decoding: The self.hidden_states buffer in CudaGraphRunner and its associated logic have been removed. Clarification is needed on how hidden states for speculative decoding are now managed within the CUDA graph context.
• EagleVerifyInput Custom Mask Initialization: The custom_mask in EagleVerifyInput is now initialized with torch.ones instead of torch.zeros. This is a fundamental change that needs justification as it could significantly alter behavior if this mask is used in attention or other critical computations.
• New topk Parameter in EagleVerifyInput: A new topk parameter has been added to EagleVerifyInput. Its role and impact on the speculative verification process should be clarified.
• Cloning of pp_proxy_tensors: Proxy tensors for pipeline parallelism are now cloned. This is a good safety measure but understanding the specific necessity would be beneficial.

Merge Readiness

This pull request introduces significant changes to the CUDA graph runner, particularly affecting speculative decoding and attention backend interactions. There are a couple of critical concerns (related to flashmla initialization and the EagleVerifyInput custom mask) that need to be addressed before this PR can be considered ready for merging. Additionally, clarifications on the handling of hidden states and the new topk parameter would be beneficial. I am unable to approve this pull request; please ensure these issues are resolved and further review is conducted before merging.