support qwen2 running on ascend npu device

[zhuyijie88]

Motivation

support qwen2 LLM running on the HUAWEI ascend npu chips.

Modifications

add is_npu() function, add npu compiler_config for torch.compile

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 @zhuyijie88, 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. This PR introduces support for running SGLang on Huawei Ascend NPU devices. The core changes involve adding NPU-specific device detection and utility functions, defining a compiler configuration for NPU using torch.compile, and integrating NPU checks into various core components and layers to ensure compatibility and proper kernel/fallback usage on the new hardware platform. This work is foundational for enabling models like Qwen2 to run efficiently on Ascend NPUs.

Highlights

• NPU Device Support: Adds the capability to detect and utilize Huawei Ascend NPU devices within SGLang.
• NPU Compiler Configuration: Introduces a specific configuration for torch.compile tailored for NPU devices, including settings like frozen_parameter, tiling_schedule_optimize, and topology_sorting_strategy.
• Integration into Core Components: Modifies various layers and utilities (custom ops, activation, layernorm, quantization, rotary embedding, device utilities) to include checks for NPU availability and conditionally use appropriate kernels or fallbacks.
• NPU System Utilities: Adds functions to query NPU memory capacity and updates the general device and compiler backend detection logic to recognize and handle NPU.

Changelog

Click here to see the changelog

• python/sglang/srt/_custom_ops.py
  • Imported the new is_npu utility function (line 7).
  • Added a check for NPU availability (not is_hip() and not is_npu()) to the condition determining whether to use custom allreduce operations (line 28).
• python/sglang/srt/configs/device_config.py
  • Added a new class method get_npu_compiler_config which returns a dictionary containing specific compiler settings for NPU devices (lines 19-26).
• python/sglang/srt/layers/activation.py
  • Imported the new is_npu utility function (line 31).
  • Added a flag _is_npu to store the result of is_npu() (line 34).
  • Added a check for NPU availability (not _is_cuda and not _is_npu) to the condition that logs a warning if sgl-kernel is not available (line 169).
• python/sglang/srt/layers/layernorm.py
  • Imported the new is_npu utility function (line 23).
  • Added a flag _is_npu to store the result of is_npu() (line 27).
  • Added a check for NPU availability (not (_is_cuda or _is_hip or _is_npu)) to the condition that logs a warning if sgl-kernel layernorm is not available (line 191).
• python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
  • Imported the new is_npu utility function (line 20).
  • Added a flag _is_npu to store the result of is_npu() (line 23).
  • Added a check for NPU availability (not _is_cuda and not _is_npu) to the condition for importing vllm custom ops (line 25).
• python/sglang/srt/layers/quantization/fp8.py
  • Imported the new is_npu utility function (line 70).
  • Added a flag _is_npu to store the result of is_npu() (line 78).
  • Added a check for NPU availability (not _is_cuda and not _is_npu) to the condition for importing vllm custom ops (line 90).
• python/sglang/srt/layers/quantization/utils.py
  • Imported the new is_npu utility function (line 9).
  • Added a flag _is_npu to store the result of is_npu() (line 12).
  • Added a check for NPU availability (not _is_cuda and not _is_npu) to the condition for importing vllm custom ops (line 14).
• python/sglang/srt/layers/rotary_embedding.py
  • Imported the new is_npu utility function (line 11).
  • Added a flag _is_npu to store the result of is_npu() (line 15).
  • Added a check for NPU availability (not (_is_cuda or _is_npu)) to the condition for using the vllm rotary embedding fallback (line 88).
• python/sglang/srt/utils.py
  • Added a new function get_npu_memory_capacity to retrieve NPU memory information using torch_npu (lines 1280-1286).
  • Updated get_device_memory_capacity to include the 'npu' case, calling the new get_npu_memory_capacity function (lines 1296-1297).
  • Updated get_device to return 'npu' or 'npu:device_id' if torch.npu.is_available() (lines 1418-1421).
  • Updated get_compiler_backend to configure and return the 'npu_backend' using torchair and the NPU compiler configuration from DeviceConfig (lines 1501-1514).

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 support for running Qwen2 models on HUAWEI Ascend NPUs, which is a valuable addition. The changes primarily involve adding utility functions for NPU detection, memory capacity retrieval, and configuring the torch.compile backend for NPUs using torchair.

The code modifications are generally clear and follow existing patterns in the repository. However, there are a few areas that could be improved for robustness and maintainability, particularly around the integration with torchair.

Crucially, the PR checklist indicates that unit tests and documentation updates are still pending. For new hardware support like this, comprehensive testing and clear documentation are essential for ensuring correctness, usability, and long-term maintainability. These should be addressed before merging.

Summary of Findings

• Missing Unit Tests: The PR checklist indicates that unit tests have not been added. For new hardware support (NPU), comprehensive unit tests are critical to verify correctness, ensure compatibility with different NPU models/drivers, and prevent regressions. This is a high-priority item.
• Missing Documentation: Documentation updates are also pending according to the checklist. Users will need guidance on prerequisites (e.g., torch_npu, torchair versions, driver requirements), setup, and any NPU-specific considerations. This is important for usability.
• Robustness of torchair Integration: In python/sglang/srt/utils.py, the integration with torchair could be made more robust by handling potential ImportError for torchair more gracefully and by clarifying if using setattr for experimental_config is the recommended approach by torchair.
• Obscure Alias for torchair: In python/sglang/srt/utils.py, torchair is imported as tng. This alias is a bit obscure. Using torch_air or simply torchair might improve readability if it doesn't cause naming conflicts. (Note: Not commented due to severity filter)

Merge Readiness

This PR makes good progress in adding NPU support. However, before it can be considered ready for merging, I strongly recommend addressing the following:

1. Add Comprehensive Unit Tests: This is crucial to ensure the NPU integration works as expected and to catch any future regressions.
2. Update Documentation: Provide clear instructions for users on how to leverage this new NPU support.
3. Address Feedback on torchair Integration: Consider the suggestions for improving the robustness of the torchair integration in utils.py.

Given the missing tests (high severity) and the medium severity comments, I am requesting changes. I am unable to approve this pull request myself; please ensure other reviewers approve it after these points are addressed.

[zhuyijie88]

@zhaochenyang20 please review this PR

[zhuyijie88]

@zhyncs @Ying1123 please review this PR.

[zhuyijie88]

/gemini review

[gemini-code-assist]

Code Review

This PR adds initial support for Huawei Ascend NPUs. The changes involve NPU detection, memory capacity retrieval, and integration with torch_air for compilation.

Key areas for improvement before merging include:

• Adding comprehensive unit and integration tests to validate NPU functionality.
• Updating documentation with setup instructions, dependencies, and NPU-specific considerations.
• Providing benchmark results to demonstrate performance characteristics on NPUs.

I've also noted a potential critical issue in rotary_embedding.py regarding the conditional import of vllm_rotary_embedding for NPUs and a query about _custom_ops.py behavior on NPUs. Please review these carefully.