[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:
@@ -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 "
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user