update metax platform deploy (#690)

Author: Wq-dd

lightx2v_platform/base/__init__.py

@@ -3,8 +3,8 @@
 from lightx2v_platform.base.ascend_npu import NpuDevice
 from lightx2v_platform.base.cambricon_mlu import MluDevice
 from lightx2v_platform.base.hygon_dcu import HygonDcuDevice
-from lightx2v_platform.base.metax import MetaxDevice
 from lightx2v_platform.base.mthreads_musa import MusaDevice
+from lightx2v_platform.base.metax_cuda import MetaxDevice
 from lightx2v_platform.base.nvidia import CudaDevice
 
 __all__ = [

lightx2v_platform/base/metax.py

@@ -0,0 +1,32 @@
+import torch
+import torch.distributed as dist
+
+from lightx2v_platform.base.nvidia import CudaDevice
+from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER
+
+
+@PLATFORM_DEVICE_REGISTER("metax_cuda")
+class MetaxDevice(CudaDevice):
+    name = "cuda"
+
+    @staticmethod
+    def init_device_env():
+        pass
+
+    @staticmethod
+    def is_available() -> bool:
+        try:
+            import torch
+
+            return torch.cuda.is_available()
+        except ImportError:
+            return False
+
+    @staticmethod
+    def get_device() -> str:
+        return "cuda"
+
+    @staticmethod
+    def init_parallel_env():
+        dist.init_process_group(backend="nccl")
+        torch.npu.set_device(dist.get_rank())

lightx2v_platform/base/metax_cuda.py

@@ -13,3 +13,5 @@
 elif PLATFORM == "ascend_npu":
     from .attn.ascend_npu import *
     from .mm.ascend_npu import *
+elif PLATFORM == "metax_cuda":
+    from .attn.metax_cuda import *

lightx2v_platform/ops/__init__.py

@@ -0,0 +1 @@
+from .metax_sage_attn2 import *

lightx2v_platform/ops/attn/metax_cuda/__init__.py

@@ -0,0 +1,45 @@
+from loguru import logger
+
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
+from lightx2v_platform.ops.attn.template import AttnWeightTemplate
+
+try:
+    from sageattention import sageattn
+except ImportError:
+    logger.info("sageattn not found, please install sageattention first")
+    sageattn = None
+
+
+@ATTN_WEIGHT_REGISTER("metax_sage_attn2")
+class MetaxSageAttn2Weight(AttnWeightTemplate):
+    def __init__(self):
+        self.config = {}
+
+    def apply(
+        self,
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_kv=None,
+        max_seqlen_q=None,
+        max_seqlen_kv=None,
+        model_cls=None,
+    ):
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+        if len(q.shape) == 3:
+            bs = 1
+            q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)
+        elif len(q.shape) == 4:
+            bs = q.shape[0]
+        x = (
+            sageattn(
+                q,
+                k,
+                v,
+                tensor_layout="NHD",
+            )[0]
+            .view(bs * max_seqlen_q, -1)
+            .type(q.dtype)
+        )
+        return x

lightx2v_platform/ops/attn/metax_cuda/metax_sage_attn2.py

@@ -6,7 +6,7 @@
 lightx2v_path=
 model_path=
 
-export PLATFORM=metax
+export PLATFORM="metax_cuda"
 export CUDA_VISIBLE_DEVICES=0
 
 # set environment variables

scripts/platforms/metax/qwen_image_i2i_2511.sh

@@ -1,13 +1,12 @@
 #!/bin/bash
 
-# System management interface: mx-smi
-
 # set path and first
-lightx2v_path=
-model_path=
+lightx2v_path=/path/to/LightX2v
+model_path=/path/to/model
 
-export PLATFORM=metax
-export CUDA_VISIBLE_DEVICES=0
+# export CUDA_VISIBLE_DEVICES=5
+export PLATFORM="metax_cuda"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
 # set environment variables
 source ${lightx2v_path}/scripts/base/base.sh