[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

@@ -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 "

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,