diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py index 598e7be9b..671adb549 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py @@ -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 " diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 5133759cb..deac46cf0 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -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, diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index 10c5841c8..87bdf6341 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/python/sglang/multimodal_gen/utils.py @@ -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)