[feature] Ascend NPU graph support by VDV1985 · Pull Request #9399 · sgl-project/sglang

VDV1985 · 2025-08-20T13:57:22Z

Motivation

NPUGraph can give significant improvement of perfrormance by reducing overhead of operators calling
Closes #8030
Second attempt to enable NPU graph functionality, first one #8027 was reverted

Modifications

Added NPU graph runner, that inherited from CudaGraphRunner. It is customers request to preserve CudaGraphRunner and not use base class.
cuda_graph_runner.py history preserved.
Added attention operation that supports by NPUGraph.update()
NPU graph runner uses the same SGLang server options as CUDA graph with combination of 'device="npu"', i.e. "disable_cuda_graph" and enabled by default.
Tests for graph added

NOTICE: it's supposed to use "export STREAMS_PER_DEVICE=32" to enable npugraph now. and you won't need it after we release the new version of torch_npu

Accuracy Tests

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[feature] accelerate ACLGraph with parallel processing of input taskupdate

gemini-code-assist

Summary of Changes

Hello @VDV1985, 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!

This pull request introduces comprehensive support for Ascend NPU graph execution, aiming to significantly enhance performance by reducing operator call overhead. It achieves this by adding a dedicated NPU graph runner that leverages existing CUDA graph infrastructure, generalizing device-specific operations across the codebase, and integrating NPU graph capabilities into the attention backend. The changes allow the system to dynamically utilize NPU graphs for optimized inference on Ascend hardware, building upon a previous attempt to enable this functionality.

Highlights

NPU Graph Runner Implementation: Introduced a new NPUGraphRunner class, specifically designed for Huawei Ascend NPUs, which inherits from the existing CudaGraphRunner to maintain a consistent interface while providing NPU-specific optimizations.
Device Agnostic Abstraction: Generalized the underlying device-specific operations (e.g., stream management, graph capture, memory pooling) within parallel_state.py and cuda_graph_runner.py to support both CUDA and NPU devices dynamically.
Attention Backend Integration: Modified the attention backend (ascend_backend.py) to integrate NPU graph capture and replay mechanisms, including new methods for managing graph-specific forward metadata and utilizing NPU-optimized fused attention operations.
Dynamic Graph Runner Selection: Updated the ModelRunner to intelligently select and initialize either the CudaGraphRunner or the new NPUGraphRunner based on the detected hardware, ensuring optimal performance for the respective device.
New NPU Graph Tests: Added new unit tests (test_ascend_graph_tp1_bf16.py, test_ascend_graph_tp2_bf16.py) to validate the functionality and performance of NPU graph support across different tensor parallelism configurations and precision settings.

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 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.

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.

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 support for Ascend NPU graphs, which is a significant feature for performance improvement. The changes primarily involve generalizing the existing CUDA graph runner to be device-agnostic and implementing the NPU-specific logic in a new NPUGraphRunner class that inherits from CudaGraphRunner. The overall approach is sound. However, I've identified a critical issue with a duplicated method definition that needs to be resolved. Additionally, there are a few opportunities for refactoring to reduce code duplication and improve maintainability, as well as a missing import that could cause a runtime error.

Alcanderian

This commit satisfies the requirements of minimizing code changes and not affecting history and downstream forks, and can now be approved

wangkeya · 2025-08-21T09:45:03Z

when i test Qwen3-32B input 8k， output500, concurreny is 50 , SGLANG will be hang up
plus

VDV1985 · 2025-08-21T10:10:25Z

when i test Qwen3-32B input 8k， output500, concurreny is 50 , SGLANG will be hang up thus

@wangkeya could you please share "npu-smi info" output and exact test parameters, i.e. sglang.launch_server params and sglang.bench_serving params (or Engine params if you use other type of testing)?

wangkeya · 2025-08-21T10:32:04Z

when i test Qwen3-32B input 8k， output500, concurreny is 50 , SGLANG will be hang up thus

@wangkeya could you please share "npu-smi info" output and exact test parameters, i.e. sglang.launch_server params and sglang.bench_serving params (or Engine params if you use other type of testing)?
sglang launch_server

bench

npu-smi info

all threads wait for result

VDV1985 · 2025-08-21T12:54:28Z

@wangkeya could you please try to set torch_npu environment variable "export STREAMS_PER_DEVICE=32" before running server? Also if you experience issues with memory allocation you can try to increase mem-fraction-static to 0.8 or 0.9.

itboyljm · 2025-08-21T13:04:12Z

Does this MR's merged code support running the DeepSeek model in ACLGraph mode?

ping1jing2 · 2025-08-21T13:29:30Z

Does this MR's merged code support running the DeepSeek model in ACLGraph mode?

I'm sorry to tell you that it's not supported yet. However, we will submit a new PR by 30/08. Thanks for your attention.

wangkeya · 2025-08-22T01:06:36Z

@wangkeya could you please try to set torch_npu environment variable "export STREAMS_PER_DEVICE=32" before running server? Also if you experience issues with memory allocation you can try to increase mem-fraction-static to 0.8 or 0.9.

thank you very much ,it works

wangkeya · 2025-09-05T10:29:14Z

@VDV1985 hello, when i test qwen3-235b-a22b, occure this error

VDV1985 · 2025-09-05T11:10:17Z

@wangkeya could you please send the error log and server/bench command lines? I suppose you have something like tensor.to("cpu") that causes stream synchronize command that is not supported by graph capturing. Also we have limited support of MOE models with graph at the moment and working on full support.

VDV1985 · 2025-09-05T11:15:56Z

@wangkeya , also you can look at this PR #10013. May be it can help.

wangkeya · 2025-09-05T12:16:51Z

#10013

i will try this pr 10013

wangkeya · 2025-09-05T14:17:29Z

#10013

i will try this pr 10013

@VDV1985 2*910B, when request concurrency is 7, sglang will be hang up。export STREAMS_PER_DEVICE=32 already set

ping1jing2 · 2025-09-05T14:47:13Z

#10013

i will try this pr 10013

@VDV1985 2*910B, when request concurrency is 7, sglang will be hang up。export STREAMS_PER_DEVICE=32 already set

@wangkeya please open a new issue and assign me, let us help you to solve the problem.

wangkeya · 2025-09-06T00:38:59Z

#10013

i will try this pr 10013

@VDV1985 2*910B, when request concurrency is 7, sglang will be hang up。export STREAMS_PER_DEVICE=32 already set

@wangkeya please open a new issue and assign me, let us help you to solve the problem.

thank you , new issue is #10091

Co-authored-by: ronnie_zheng <zl19940307@163.com> Co-authored-by: yezhifeng (D) <y00897525@china.huawei.com> Co-authored-by: anon189Ty <Stari_Falcon@outlook.com> Co-authored-by: Maksim <makcum888e@mail.ru> Co-authored-by: ssshinigami <44640852+ssshinigami@users.noreply.github.com>

VDV1985 and others added 30 commits July 14, 2025 17:33

NPU graph support

485515f

clean code

cf708c0

add missing function

729ddda

clean code

9bfee5b

fix for prev commit

f0cded4

remap cuda.Stream to npu.Stream

7c42b4c

fix typo

a36f9f8

substitute cuda with npu only for npu

8b7c373

refactor NPUGraphRunner

9297f55

fix missing import

6eb161e

fix missing parameter

07e4c47

fix typos

d07594d

fix typos

a7d16e0

Merge branch 'sgl-project:main' into npu_graph_support

8cdc791

Merge branch 'sgl-project:main' into npu_graph_support

0e2b261

merge with upstream

f20944f

minor attn backend update, move CI to vevelopers version of torch_npu

d7bdffa

Merge branch 'main' into npu_graph_support

e86319f

update torch-npu version to release

9142b2b

add thread

3004107

add thread

76c1bd6

add thread

2bec597

refactor

1e57f56

fix

6ee14e1

Merge pull request #2 from anon189Ty/npu_graph_support

3ac5849

[feature] accelerate ACLGraph with parallel processing of input taskupdate

Merge remote-tracking branch 'base/main' into npu_graph_support

df3ba1d

fix

465f00d

fix

10b78e1

fix comments

5777fd0

graph tests

e7c3326

VDV1985 requested review from BBuf, Edwardf0t1, HaiShaw, ch-wan, kushanam and yizhang2077 as code owners August 20, 2025 13:57

gemini-code-assist Bot reviewed Aug 20, 2025

View reviewed changes

Comment thread python/sglang/srt/layers/attention/ascend_backend.py Outdated

Comment thread python/sglang/srt/layers/attention/ascend_backend.py

Comment thread python/sglang/srt/model_executor/npu_graph_runner.py

Comment thread python/sglang/srt/model_executor/npu_graph_runner.py Outdated

fix gemini comments

64804a8

ping1jing2 assigned Alcanderian Aug 20, 2025

Alcanderian approved these changes Aug 21, 2025

View reviewed changes

zhyncs merged commit 2c4b4b7 into sgl-project:main Aug 21, 2025
116 of 122 checks passed

iforgetmyname mentioned this pull request Aug 27, 2025

[Feature] Support NPUGraph for DeepSeek on Ascend NPU #9355

Merged

4 tasks

ping1jing2 deleted the npu_graph_support_redesign branch May 4, 2026 07:22

Conversation

VDV1985 commented Aug 20, 2025 • edited by ping1jing2 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Alcanderian left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

wangkeya commented Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VDV1985 commented Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wangkeya commented Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VDV1985 commented Aug 21, 2025

Uh oh!

itboyljm commented Aug 21, 2025

Uh oh!

ping1jing2 commented Aug 21, 2025

Uh oh!

wangkeya commented Aug 22, 2025

Uh oh!

wangkeya commented Sep 5, 2025

Uh oh!

VDV1985 commented Sep 5, 2025

Uh oh!

VDV1985 commented Sep 5, 2025

Uh oh!

wangkeya commented Sep 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wangkeya commented Sep 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ping1jing2 commented Sep 5, 2025

Uh oh!

wangkeya commented Sep 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

VDV1985 commented Aug 20, 2025 •

edited by ping1jing2

Loading

wangkeya commented Aug 21, 2025 •

edited

Loading

VDV1985 commented Aug 21, 2025 •

edited

Loading

wangkeya commented Aug 21, 2025 •

edited

Loading

wangkeya commented Sep 5, 2025 •

edited

Loading

wangkeya commented Sep 5, 2025 •

edited

Loading

wangkeya commented Sep 6, 2025 •

edited

Loading