Skip to content

[AMD] Improve openai/gpt-oss performance#21020

Merged
merrymercy merged 5 commits intosgl-project:mainfrom
HaiShaw:optimize-schedule
Mar 21, 2026
Merged

[AMD] Improve openai/gpt-oss performance#21020
merrymercy merged 5 commits intosgl-project:mainfrom
HaiShaw:optimize-schedule

Conversation

@kkHuang-amd
Copy link
Copy Markdown
Collaborator

@kkHuang-amd kkHuang-amd commented Mar 20, 2026 •
edited
Loading

Motivation

Check the torch trace, found some blocking call (gpu memory -> cpu memory) to block host side to dispatch workload to GPU, lead some bubble of GPU run between steps.
Non-optimzed
image

optimzed
image

Modifications

Original call forward_batch.seq_lens.max().item()
Change call forward_batch.seq_lens_cpu.max().item()

Original call forward_batch.extend_seq_lens.max().item()
Change call max(forward_batch.extend_seq_lens_cpu)

Accuracy Tests

Benchmarking and Profiling

image

Checklist

Review Process

  1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
  4. After green CI and required approvals, ask Merge Oncalls to merge.

@gemini-code-assist
Copy link
Copy Markdown
Contributor

gemini-code-assist Bot commented Mar 20, 2026

Summary of Changes

Hello, I'm Gemini Code Assist1! 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!

This pull request significantly improves the performance of the openai/gpt-oss model by optimizing how sequence length metadata is handled. The changes focus on reducing unnecessary GPU-to-CPU data transfers, which were identified as a source of performance bubbles. By performing maximum length calculations on CPU tensors and passing CPU-resident data directly, the host side can dispatch workloads to the GPU more efficiently, leading to a smoother execution flow and better overall throughput.

Highlights

  • Performance Optimization: Addressed performance bottlenecks caused by blocking GPU-to-CPU memory transfers, specifically identified through torch trace analysis, to reduce GPU idle time between steps.
  • CPU-based Max Length Calculation: Modified the calculation of max_kv_len and max_q_len to utilize CPU-resident tensors (seq_lens_cpu, extend_seq_lens_cpu) where possible, minimizing costly GPU-to-CPU synchronizations.
  • Refactored Metadata Initialization: Streamlined the init_forward_metadata and init_forward_metadata_replay_cuda_graph functions by passing seq_lens_cpu directly and using pre-calculated max_kv_len values, enhancing efficiency.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

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[bot]
gemini-code-assist Bot reviewed Mar 20, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request aims to improve performance by removing blocking GPU-to-CPU memory transfers. The changes replace .item() calls on GPU tensors with computations on pre-existing CPU tensors, which should reduce synchronization overhead. Additionally, the changes correct a potential bug by using the maximum sequence length instead of the total number of tokens in certain calculations. While the overall approach is sound, I've found one critical issue introduced by a typo that needs to be fixed.

Comment thread python/sglang/srt/layers/attention/aiter_backend.py Outdated
@kkHuang-amd kkHuang-amd changed the title Improve openai/gpt-oss performance [AMD] Improve openai/gpt-oss performance Mar 20, 2026
hubertlu-tw
hubertlu-tw approved these changes Mar 20, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@hubertlu-tw hubertlu-tw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Nice catch!

HaiShaw
HaiShaw approved these changes Mar 21, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@HaiShaw HaiShaw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed incorrect modification

@merrymercy merrymercy merged commit 3f0ba02 into sgl-project:main Mar 21, 2026
53 of 81 checks passed
0-693 pushed a commit to 0-693/sglang that referenced this pull request Mar 25, 2026
@kkHuang-amd @hubertlu-tw
@HaiShaw @0-693
[AMD] Improve openai/gpt-oss performance (sgl-project#21020)
426c883 
Co-authored-by: root <root@smci355-ccs-aus-m15-21.cs-aus.dcgpu>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
dutsc pushed a commit to dutsc/sglang that referenced this pull request Mar 30, 2026
@kkHuang-amd @hubertlu-tw @HaiShaw
[AMD] Improve openai/gpt-oss performance (sgl-project#21020)
76013e8 
Co-authored-by: root <root@smci355-ccs-aus-m15-21.cs-aus.dcgpu>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026
@kkHuang-amd @hubertlu-tw @HaiShaw
[AMD] Improve openai/gpt-oss performance (sgl-project#21020)
162f3a2 
Co-authored-by: root <root@smci355-ccs-aus-m15-21.cs-aus.dcgpu>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
yhyang201 pushed a commit to yhyang201/sglang that referenced this pull request Apr 22, 2026
@kkHuang-amd @hubertlu-tw @HaiShaw
[AMD] Improve openai/gpt-oss performance (sgl-project#21020)
b18db5b 
Co-authored-by: root <root@smci355-ccs-aus-m15-21.cs-aus.dcgpu>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@HaiShaw HaiShaw HaiShaw approved these changes

@hubertlu-tw hubertlu-tw hubertlu-tw approved these changes

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@Qiaolin-Yu Qiaolin-Yu Awaiting requested review from Qiaolin-Yu Qiaolin-Yu is a code owner

@hebiao064 hebiao064 Awaiting requested review from hebiao064 hebiao064 is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

run-ci

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@kkHuang-amd @HaiShaw @hubertlu-tw @merrymercy