Extend cuda graph capture bs for B200

[Fridge003]

Motivation

The extension of cuda_graph bs on B200 was deleted in #6803
This PR adds it back

cc @merrymercy

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 @Fridge003, 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 with a summary of this pull request. This PR, authored by Fridge003, aims to re-introduce the extension of CUDA graph capture batch sizes specifically for B200 GPUs. According to the description, this extension was previously removed in pull request #6803 and is now being added back. The change involves modifying the logic that determines the batch sizes to capture for CUDA graphs based on available GPU memory.

Highlights

• CUDA Graph Capture: The core of this PR is related to configuring CUDA graph capture, a technique used to optimize performance by capturing sequences of GPU operations.
• B200 GPU Support: This change specifically targets high-memory GPUs like the NVIDIA B200, enabling larger batch sizes for CUDA graph capture on these devices.
• Batch Size Extension: The PR adds a condition to extend the range of batch sizes captured for CUDA graphs when the detected GPU memory exceeds 180GB.

Changelog

• python/sglang/srt/model_executor/cuda_graph_runner.py
  • Added a new conditional block to get_batch_sizes_to_capture (around line 141) that checks if the GPU memory is greater than 180GB.
  • If the GPU memory is greater than 180GB, the list of batch sizes to capture (capture_bs) is extended with values ranging from 256 to 320, in steps of 16.

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 reintroduces an extension for CUDA graph capture batch sizes, specifically targeting B200 GPUs. The change is minimal and directly addresses the issue described. The code is clear and easy to understand. One point of attention is the consistency of units used for GPU memory thresholds.

Summary of Findings

• Inconsistent GPU Memory Unit: The new GPU memory threshold 180 * 1000 uses a different multiplier (1000) for 'K' compared to the existing threshold 96 * 1024 (which uses 1024). Since get_device_memory_capacity() returns values in MiB, using 1024 consistently for such thresholds (e.g., 180 * 1024 for 180 GiB) would improve clarity and ensure consistent unit interpretation.

Merge Readiness

This PR is well-focused on restoring functionality for B200 GPUs. There is one medium-severity suggestion regarding the consistency of units for GPU memory thresholds that should be considered. Addressing or clarifying this point would be beneficial. As a language model, I am not authorized to approve pull requests; this decision should be made by the designated reviewers after considering the feedback.