[diffusion] hardware: support diffusion models on MTGPU (multi-GPU, 5/N) (#17318)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-02-04 04:44:22 +08:00
committed by GitHub
parent acf724b036
commit ec2461bc16
3 changed files with 9 additions and 5 deletions

View File

@@ -52,8 +52,8 @@ def find_nccl_library() -> str:
"""
We either use the library file specified by the `VLLM_NCCL_SO_PATH`
environment variable, or we find the library file brought by PyTorch.
After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
found by `ctypes` automatically.
After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2`
can be found by `ctypes` automatically.
"""
so_file = envs.SGLANG_DIFFUSION_NCCL_SO_PATH
@@ -68,8 +68,10 @@ def find_nccl_library() -> str:
so_file = "libnccl.so.2"
elif torch.version.hip is not None:
so_file = "librccl.so.1"
elif hasattr(torch.version, "musa") and torch.version.musa is not None:
so_file = "libmccl.so.2"
else:
raise ValueError("NCCL only supports CUDA and ROCm backends.")
raise ValueError("NCCL only supports CUDA, ROCm and MUSA backends.")
logger.info("Found nccl from library %s", so_file)
return str(so_file)