minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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-03 12: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
python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py
+1 -1
View File
@@ -277,7 +277,7 @@ class NCCLLibrary:
except Exception as e: except Exception as e:
logger.error( logger.error(
"Failed to load NCCL library from %s ." "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 " "Otherwise, the nccl library might not exist, be corrupted "
"or it does not support the current platform %s." "or it does not support the current platform %s."
"If you already have the library, please set the " "If you already have the library, please set the "
python/sglang/multimodal_gen/runtime/models/encoders/clip.py
+3 -1
View File
@@ -230,10 +230,12 @@ class CLIPAttention(nn.Module):
key_states = key_states.transpose(1, 2) key_states = key_states.transpose(1, 2)
value_states = value_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. # ROCm: Using both is_causal=True and attn_mask causes NaN.
# Use is_causal=True alone (padding mask not needed for CLIP # Use is_causal=True alone (padding mask not needed for CLIP
# since pooler_output comes from EOS token before padding). # 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( attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states, query_states,
key_states, key_states,
python/sglang/multimodal_gen/utils.py
+5 -3
View File
@@ -52,8 +52,8 @@ def find_nccl_library() -> str:
""" """
We either use the library file specified by the `VLLM_NCCL_SO_PATH` We either use the library file specified by the `VLLM_NCCL_SO_PATH`
environment variable, or we find the library file brought by PyTorch. environment variable, or we find the library file brought by PyTorch.
After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2`
found by `ctypes` automatically. can be found by `ctypes` automatically.
""" """
so_file = envs.SGLANG_DIFFUSION_NCCL_SO_PATH so_file = envs.SGLANG_DIFFUSION_NCCL_SO_PATH
@@ -68,8 +68,10 @@ def find_nccl_library() -> str:
so_file = "libnccl.so.2" so_file = "libnccl.so.2"
elif torch.version.hip is not None: elif torch.version.hip is not None:
so_file = "librccl.so.1" so_file = "librccl.so.1"
elif hasattr(torch.version, "musa") and torch.version.musa is not None:
so_file = "libmccl.so.2"
else: 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) logger.info("Found nccl from library %s", so_file)
return str(so_file) return str(so_file)