[XPU][Diffusion] Enable MiniMax H3 on XPU platforms (#33366)
This commit is contained in:
@@ -119,3 +119,30 @@ class XPUAttentionImpl(AttentionImpl):
|
|||||||
|
|
||||||
result = out.reshape(bsz, seqlen_q, nheads_q, d)
|
result = out.reshape(bsz, seqlen_q, nheads_q, d)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
def forward_varlen(
|
||||||
|
self,
|
||||||
|
query: torch.Tensor,
|
||||||
|
key: torch.Tensor,
|
||||||
|
value: torch.Tensor,
|
||||||
|
*,
|
||||||
|
cu_seqlens: torch.Tensor,
|
||||||
|
max_seqlen: int,
|
||||||
|
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
del cu_seqlens_host
|
||||||
|
q_ = query.contiguous()
|
||||||
|
k_ = key.contiguous()
|
||||||
|
v_ = value.contiguous()
|
||||||
|
output = flash_attn_func(
|
||||||
|
q=q_,
|
||||||
|
k=k_,
|
||||||
|
v=v_,
|
||||||
|
cu_seqlens_q=cu_seqlens,
|
||||||
|
cu_seqlens_k=cu_seqlens,
|
||||||
|
max_seqlen_q=max_seqlen,
|
||||||
|
max_seqlen_k=max_seqlen,
|
||||||
|
softmax_scale=self.softmax_scale,
|
||||||
|
causal=self.causal,
|
||||||
|
)
|
||||||
|
return output[0] if isinstance(output, tuple) else output
|
||||||
|
|||||||
@@ -926,6 +926,7 @@ class ComponentResidencyManager:
|
|||||||
current_platform.is_cuda()
|
current_platform.is_cuda()
|
||||||
or current_platform.is_rocm()
|
or current_platform.is_rocm()
|
||||||
or current_platform.is_npu()
|
or current_platform.is_npu()
|
||||||
|
or current_platform.is_xpu()
|
||||||
)
|
)
|
||||||
return is_supported_platform and current_platform.is_device_type(
|
return is_supported_platform and current_platform.is_device_type(
|
||||||
self._module_device(module)
|
self._module_device(module)
|
||||||
|
|||||||
+1
@@ -48,6 +48,7 @@ def minimax_h3_scoped_encode_rng(seed: int, device: torch.device | None = None):
|
|||||||
current_platform.is_cuda()
|
current_platform.is_cuda()
|
||||||
or current_platform.is_rocm()
|
or current_platform.is_rocm()
|
||||||
or current_platform.is_npu()
|
or current_platform.is_npu()
|
||||||
|
or current_platform.is_xpu()
|
||||||
)
|
)
|
||||||
if (
|
if (
|
||||||
device is not None
|
device is not None
|
||||||
|
|||||||
+2
-1
@@ -679,9 +679,10 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
|||||||
or current_platform.is_cpu()
|
or current_platform.is_cpu()
|
||||||
or current_platform.is_mps()
|
or current_platform.is_mps()
|
||||||
or current_platform.is_npu()
|
or current_platform.is_npu()
|
||||||
|
or current_platform.is_xpu()
|
||||||
):
|
):
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
"MiniMax H3 full-loop denoise requires CPU, CUDA, MPS, or Ascend NPU"
|
"MiniMax H3 full-loop denoise requires CPU, CUDA, MPS, XPU, or Ascend NPU"
|
||||||
)
|
)
|
||||||
|
|
||||||
device = current_platform.get_local_torch_device()
|
device = current_platform.get_local_torch_device()
|
||||||
|
|||||||
Reference in New Issue
Block a user