Remove unnecessary kernels of num_token_non_padded

[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! Gemini here, providing a summary of this pull request to help everyone quickly understand the changes. This PR, titled "Remove unnecessary kernels of num_token_non_padded", aims to optimize the code by making the calculation and usage of the num_token_non_padded value conditional. Based on the code changes, it appears this value is only necessary when specific Mixture of Experts (MoE) features are enabled. By adding a check for these features, the PR avoids unnecessary computation and tensor allocation related to num_token_non_padded when those features are not active, potentially improving performance or reducing overhead in those cases.

Highlights

• Conditional Calculation: The calculation and creation of the num_token_non_padded tensor in ForwardBatchInfo.init_new is now conditional, only occurring if specific MoE features are enabled.
• Conditional Assignment: The assignment of the num_token_non_padded value within the CUDA graph replay preparation (CudaGraphRunner.replay_prepare) is also made conditional based on the same MoE feature check.
• New Helper Function: A new helper function enable_num_token_non_padded is introduced to encapsulate the logic for checking if the relevant MoE features are enabled via the server arguments.

Changelog

• python/sglang/srt/model_executor/cuda_graph_runner.py
  • Imported the new enable_num_token_non_padded helper function (line 38).
  • Wrapped the assignment to self.num_token_non_padded within an if enable_num_token_non_padded(...) block (lines 554-555).
• python/sglang/srt/model_executor/forward_batch_info.py
  • Introduced a local variable num_token_non_padded initialized to None (line 276).
  • Conditionally calculated and assigned the tensor value to num_token_non_padded only if enable_num_token_non_padded returns true (lines 277-280).
  • Passed the local num_token_non_padded variable (which might be None) to the class constructor (line 310).
  • Added the new helper function enable_num_token_non_padded which checks for enable_ep_moe or enable_deepep_moe in server arguments (lines 614-615).

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 aims to optimize the handling of num_token_non_padded by making its computation and usage conditional, primarily based on whether Mixture of Experts (MoE) features like enable_ep_moe or enable_deepep_moe are active. This is a good step towards reducing unnecessary computations when these features are not in use.

The changes involve conditionally initializing ForwardBatch.num_token_non_padded to None and conditionally updating CudaGraphRunner.self.num_token_non_padded during CUDA graph replay. While the intent is clear, there are a couple of points that need clarification to ensure correctness, particularly regarding the handling of None values by downstream code and the potential for stale data in CUDA graph replays.

Summary of Findings

• Potential Stale Data in CUDA Graph Replay: In cuda_graph_runner.py, self.num_token_non_padded (a CUDA graph input) is not updated if MoE is disabled during replay. This could lead to the graph using stale data if kernels within the graph depend on this value and are not properly gated by MoE flags. This is flagged as a high-severity concern.
• Handling of None for ForwardBatch.num_token_non_padded: In forward_batch_info.py, ForwardBatch.num_token_non_padded can now be None. It's important to ensure all downstream code that uses this attribute can correctly handle None or is only active when MoE features are enabled. This is flagged as a medium-severity concern.
• Docstring for enable_num_token_non_padded: Consider adding a docstring to the new enable_num_token_non_padded function in forward_batch_info.py to explain its purpose, what num_token_non_padded is used for, and why its usage is tied to MoE. This would improve code clarity for future maintainers. (Severity: low, not commented inline due to settings)

Merge Readiness

The pull request introduces a sensible optimization by making the handling of num_token_non_padded conditional on MoE settings. However, there are a couple of potential issues raised that need clarification or addressing before this PR can be safely merged:

1. A high-severity concern regarding the CUDA graph replay in cuda_graph_runner.py, where self.num_token_non_padded might use stale data if MoE is disabled.
2. A medium-severity concern about ensuring all consumers of ForwardBatch.num_token_non_padded can handle it being None when MoE is disabled.

I recommend addressing these points to ensure the correctness and robustness of the changes. Therefore, I am requesting changes. As an AI, I am not authorized to approve pull requests; please ensure further review and approval from other team members after addressing the feedback.