Add Ascend NPU support for xDiT

[ChenTaoyu-SJTU]

Currently, the xDiT is not support NPU platform, and we prepare to add NPU support for xDiT.

We need to complete this task in several phases:

• Adapt a specific model for single-card execution. Add NPU support for one model in single card #566
• Adapt a specific model for multi-card execution with model parallelism.
• Support multiple models running on a single card (model adaptation).
• Enable multi-card parallelism for as many models as possible.

Here is the detail illustration of above stages:

1. First, we will submit the PR for Phase 1, which adds support for running the black-forest-labs/FLUX.1-dev model on a single NPU. The main changes are in the environment and the parallel communication backend. Details can be found in the First PR: Add NPU support for one model in single card #566 .

[ChenTaoyu-SJTU]

@feifeibear Do you have any suggestions for this issue and PR?

[GuangyuZhu04]

Currently, the xDiT is not support NPU platform, and we prepare to add NPU support for xDiT.

We need to complete this task in several phases:

• Adapt a specific model for single-card execution. Add NPU support for one model in single card #566[ ] Adapt a specific model for multi-card execution with model parallelism.[ ] Support multiple models running on a single card (model adaptation).[ ] Enable multi-card parallelism for as many models as possible.

Here is the detail illustration of above stages:

1. First, we will submit the PR for Phase 1, which adds support for running the black-forest-labs/FLUX.1-dev model on a single NPU. The main changes are in the environment and the parallel communication backend. Details can be found in the First PR: Add NPU support for one model in single card #566 .

Do you have any Ascend image environment available that can currently run this code?

[ChenTaoyu-SJTU]

@GuangyuZhu04 I don't use a specific image to run this code, but using the environment after I installed the vllm_ascend and then install the xDiT, I guess the key packages should be torch=2.7.1, torch-npu==2.7.1.dev20250724, with the mentioned two packages, then install the xDiT, you may run xDiT well. Or you can just get an image from vllm_ascend, then install the xDiT which will also work.

[GuangyuZhu04]

@ChenTaoyu-SJTU

Got it! Thank you. I will try the image from vllm_ascend first.

[GuangyuZhu04]

@ChenTaoyu-SJTU I already built the vllm_ascend evironment with the installed torch and torch_npu. What shall I do next?
It seems that I cannot directly run the shell code and python code in #566 .

[ChenTaoyu-SJTU]

@GuangyuZhu04 whether you installed the xDiT from source also? And what's the concrete error when you run it?

[GuangyuZhu04]

@ChenTaoyu-SJTU How can I install xDiT from source? Since torch is installed in cpu version, I cannot use the command, 'pip install "xfuser[diffusers,flash-attn]"', to install xfuser. When I run the launch.sh, it still use torch CUDA api, and therefore, I cannot run the sd3.py.

[GuangyuZhu04]

@ChenTaoyu-SJTU This is my code error:

[root@10-155-205-228 xDiT_example]# ./launch.sh
++ torchrun --nproc_per_node=8 --start_method=spawn /workspace/xDiT_example/sd3.py --model /root/.cache/modelscope/hub/models/stabilityai/stable-diffusion-3-medium-diffusers --height 1024 --width 1024 --no_use_resolution_binning --guidance_scale 3.5 --num_inference_steps 50 --warmup_steps 1 --prompt 'brown dog laying on the ground with a metal bowl in front of him.' 'A small cat.' 'A good man' --tensor_parallel_degree 2 --data_parallel_degree 2
W1027 07:28:35.695000 822 site-packages/torch/distributed/run.py:766]
W1027 07:28:35.695000 822 site-packages/torch/distributed/run.py:766] *****************************************
W1027 07:28:35.695000 822 site-packages/torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W1027 07:28:35.695000 822 site-packages/torch/distributed/run.py:766] *****************************************
Traceback (most recent call last):
File "/workspace/xDiT_example/sd3.py", line 6, in
from xfuser import xFuserStableDiffusion3Pipeline, xFuserArgs
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/init.py", line 1, in
from xfuser.model_executor.pipelines import (
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/model_executor/pipelines/init.py", line 1, in
from .base_pipeline import xFuserPipelineBaseWrapper
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/model_executor/pipelines/base_pipeline.py", line 13, in
from xfuser.core.distributed.group_coordinator import GroupCoordinator
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/core/init.py", line 1, in
from .cache_manager import CacheManager
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/core/cache_manager/init.py", line 1, in
from .cache_manager import CacheManager
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/core/cache_manager/cache_manager.py", line 4, in
from xfuser.core.distributed.runtime_state import runtime_state_is_initialized
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/core/distributed/init.py", line 1, in
from .parallel_state import (
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/core/distributed/parallel_state.py", line 11, in
import xfuser.envs as envs
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/envs.py", line 218, in
PACKAGES_CHECKER = PackagesEnvChecker()
^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/envs.py", line 151, in new
cls._instance.initialize()
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/envs.py", line 156, in initialize
"has_flash_attn": self.check_flash_attn(),
^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/xfuser/envs.py", line 169, in check_flash_attn
gpu_name = torch.cuda.get_device_name(device)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/torch/cuda/init.py", line 544, in get_device_name
return get_device_properties(device).name
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/torch/cuda/init.py", line 576, in get_device_properties
_lazy_init() # will define _get_device_properties
^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/torch/cuda/init.py", line 363, in _lazy_init
raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled
[ERROR] 2025-10-27-07:28:47 (PID:894, Device:-1, RankID:-1) ERR99999 UNKNOWN applicaiton exception

[ChenTaoyu-SJTU]

@GuangyuZhu04 just using pip install -e . is fine. By the way, may be --nproc_per_node=4 is better when tp=2 and dp=2

[GuangyuZhu04]

@ChenTaoyu-SJTU We have successfully got the code running. Thank you so much for your assistance! 👍

[ChenTaoyu-SJTU]

@GuangyuZhu04 you are welcome. Happy to see you make the code working.

[nuclearwu]

@GuangyuZhu04 you are welcome. Happy to see you make the code working.

I have successfully executed examples/run.sh by modifying flux_example.py to adapt to the npu. So, how can I run the model on the npu in a multi-card manner? Could you share it. @ChenTaoyu-SJTU

[ChenTaoyu-SJTU]

@nuclearwu here is the example for the stable_diffusion using tp=2 and dp=2: #566, you can modify it to adapt the flux

[nuclearwu]

@nuclearwu here is the example for the stable_diffusion using tp=2 and dp=2: #566, you can modify it to adapt the flux

@ChenTaoyu-SJTU OK, have done. Are --pipefusion_parallel_degree --ulysses-degree and --ring-degree not supported for ascend npu?