chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py
View File

@@ -277,7 +277,7 @@ class NCCLLibrary:
except Exception as e:
logger.error(
"Failed to load NCCL library from %s ."
"It is expected if you are not running on NVIDIA/AMD GPUs."
"It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs."
"Otherwise, the nccl library might not exist, be corrupted "
"or it does not support the current platform %s."
"If you already have the library, please set the "

4
python/sglang/multimodal_gen/runtime/models/encoders/clip.py
View File

@@ -230,10 +230,12 @@ class CLIPAttention(nn.Module):
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
if current_platform.is_rocm():
if current_platform.is_rocm() or current_platform.is_musa():
# ROCm: Using both is_causal=True and attn_mask causes NaN.
# Use is_causal=True alone (padding mask not needed for CLIP
# since pooler_output comes from EOS token before padding).
# XXX (MUSA): Torch SDPA on MUSA currently does not support
# using both `attn_mask` and `is_causal=True` simultaneously.
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states,
key_states,

8
python/sglang/multimodal_gen/utils.py
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)