[diffusion][MUSA] fix: MUSA platform breakage caused by PR #13662

[yeahdongcn]

Motivation

PR #13662 introduced a regression that broke MUSA platform functionality. This PR fixes the issue and restores MUSA support.

Modifications

Align the MUSA implementation with the existing CUDA platform behavior.

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.
• Follow the SGLang code style guidance.

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]

Summary of Changes

Hello @yeahdongcn, 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 provides a critical fix for the MUSA platform, which experienced a regression due to a previous change. The primary goal is to re-establish MUSA functionality by ensuring its implementation correctly handles device identification and memory queries, particularly in distributed settings, thereby aligning its behavior with existing CUDA platform standards.

Highlights

• MUSA Platform Fix: This pull request addresses and resolves a regression introduced by PR [NPU][diffusion] model: support WAN/FLUX/Qwen-Image/Qwen-Image-edit on Ascend #13662, which caused breakage in the MUSA platform functionality. The changes aim to restore full MUSA support.
• Distributed Device Identification: A new class method, get_local_torch_device, has been added to MusaPlatformBase to correctly identify the local MUSA device using the LOCAL_RANK environment variable, aligning with distributed training practices.
• Distributed GPU Memory Query Correction: The get_available_gpu_memory function has been updated to correctly determine the device_id in distributed environments by utilizing torch.distributed.get_rank() when torch.distributed is initialized, ensuring accurate memory reporting for MUSA devices.

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

Changelog

• python/sglang/multimodal_gen/runtime/platforms/musa.py
  • Imported the envs module to access environment variables.
  • Added the get_local_torch_device class method to MusaPlatformBase to return a torch.device object for the local MUSA device based on envs.LOCAL_RANK.
  • Modified the get_available_gpu_memory function to set device_id to torch.distributed.get_rank() if torch.distributed is initialized, ensuring correct device targeting in distributed setups.

Activity

• The pull request was created to fix a regression in MUSA platform functionality caused by PR [NPU][diffusion] model: support WAN/FLUX/Qwen-Image/Qwen-Image-edit on Ascend #13662.
• The author, yeahdongcn, has provided a clear motivation and details on modifications made to align MUSA implementation with CUDA behavior.

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.

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

[yeahdongcn]

@Makcum888e Please also be aware of this.

[gemini-code-assist]

Code Review

This pull request aims to fix a regression on the MUSA platform by aligning its implementation with the existing CUDA platform. The changes introduce get_local_torch_device and modify get_available_gpu_memory to correctly handle device IDs in a distributed context. The changes are logical and consistent with the goal of aligning with the CUDA implementation. I have provided one suggestion to improve the robustness of device selection in multi-node distributed environments.

[mickqian]

/tag-and-rerun-ci

[Makcum888e]

Please also be aware of this.

please add CI tests for your platform to prevent similar problems in the future

[yeahdongcn]

Please also be aware of this.

please add CI tests for your platform to prevent similar problems in the future

Yes, @johnnycxm is working on that.

[yeahdongcn]

It appears that all the failing AMD CI jobs are related to the human-eval installation. Given that, could @mickqian @Kangyan-Zhou please proceed with merging this PR? Thanks!

[mickqian]

/rerun-failed-ci