feat(ascend): add MiniMax-H3 fused RoPE and MindIE SD compile backend

configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json

@@ -0,0 +1,30 @@
+{
+    "infer_steps": 30,
+    "target_video_length": 362,
+    "target_height": 768,
+    "target_width": 1344,
+    "fps": 24,
+    "enable_cfg": false,
+    "cpu_offload": false,
+    "offload_granularity": "model",
+    "text_encoder_cpu_offload": true,
+    "text_encoder_tensor_parallel": false,
+    "vae_cpu_offload": true,
+    "vae_decode_parallel": true,
+    "lazy_load": false,
+    "unload_modules": false,
+    "attn_type": "npu_flash_attn",
+    "rms_type": "npu_rms_norm",
+    "rope_type": "minimax_h3_npu_rope",
+    "feature_caching": "NoCaching",
+    "use_compile": true,
+    "compile_backend": "mindie",
+    "video_flow_shift": 12.0,
+    "audio_flow_shift": 3.0,
+    "vae_spatial_scale_factor": 16,
+    "audio_sampling_rate": 32000,
+    "audio_latents_per_second": 40,
+    "audio_channels": 2,
+    "keep_latents_dtype_in_scheduler": true,
+    "parallel": {"tensor_p_size": 1, "seq_p_size": 4, "seq_p_attn_type": "ulysses"},
+}

configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json

@@ -0,0 +1,30 @@
+{
+    "infer_steps": 30,
+    "target_video_length": 120,
+    "target_height": 768,
+    "target_width": 1344,
+    "fps": 24,
+    "enable_cfg": false,
+    "cpu_offload": false,
+    "offload_granularity": "model",
+    "text_encoder_cpu_offload": false,
+    "text_encoder_tensor_parallel": false,
+    "vae_cpu_offload": true,
+    "vae_decode_parallel": true,
+    "lazy_load": false,
+    "unload_modules": false,
+    "attn_type": "npu_flash_attn",
+    "rms_type": "npu_rms_norm",
+    "rope_type": "minimax_h3_npu_rope",
+    "feature_caching": "NoCaching",
+    "use_compile": true,
+    "compile_backend": "mindie",
+    "video_flow_shift": 12.0,
+    "audio_flow_shift": 3.0,
+    "vae_spatial_scale_factor": 16,
+    "audio_sampling_rate": 32000,
+    "audio_latents_per_second": 40,
+    "audio_channels": 2,
+    "keep_latents_dtype_in_scheduler": true,
+    "parallel": {"tensor_p_size": 1, "seq_p_size": 4, "seq_p_attn_type": "ulysses"},
+}

lightx2v/common/ops/attn/ulysses_a2a.py

@@ -33,7 +33,11 @@ class TorchUlyssesA2A:
     """Ulysses all-to-all implemented by ``torch.distributed``."""
 
     @staticmethod
+    @torch._dynamo.disable
     def exchange(input_tensor, group=None, async_op=False):
+        # torch._dynamo.disable: keep the collective out of the compiled graph.
+        # Tracing it degrades HCCL all_to_all_single to the slow variable-length
+        # alltoallv path; running eager keeps the surrounding compute fused.
         output_tensor = torch.empty_like(input_tensor)
         work = dist.all_to_all_single(output_tensor, input_tensor, group=group, async_op=async_op)
         return output_tensor, work

lightx2v/common/transformer_infer/transformer_infer.py

@@ -8,9 +8,39 @@
 class BaseTransformerInfer(ABC):
     def init_compile(self, config):
         self.use_compile = config.get("use_compile", False)
+        # compile_backend: "default" -> plain torch.compile; "mindie" -> MindieSDBackend
+        self.compile_backend = config.get("compile_backend", "default")
+        # compile_dynamic: "None" (default auto), True, or False. H3 per-step
+        # sequence is fixed, so False skips per-call shape guard checks.
+        self.compile_dynamic = config.get("compile_dynamic", None)
         self.compiled_blocks = {}
+        self._compile_backend_obj = self._create_compile_backend() if self.use_compile else None
         if self.use_compile:
-            logger.info(f"[Compile] Using torch.compile for {type(self).__name__}")
+            logger.info(f"[Compile] Using torch.compile (backend={self.compile_backend}) for {type(self).__name__}")
+
+    def _create_compile_backend(self):
+        """Instantiate the configured compile backend, or None for the default.
+
+        Reuse ONE backend instance: a fresh MindieSDBackend() per call makes
+        Dynamo see a different backend callable each time (BACKEND_MATCH
+        recompilation until recompile_limit, then silent eager fallback).
+        Unknown names and unavailable optional backends degrade to the default
+        torch.compile with a warning.
+        """
+        if self.compile_backend not in ("default", "mindie"):
+            logger.warning(f"[Compile] Unknown compile_backend={self.compile_backend!r}; expected 'default' or 'mindie'. Falling back to 'default'.")
+            self.compile_backend = "default"
+            return None
+        if self.compile_backend == "default":
+            return None
+        try:
+            from mindiesd.compilation import MindieSDBackend
+
+            return MindieSDBackend()
+        except Exception as e:  # pragma: no cover - optional dependency path
+            logger.warning(f"[Compile] MindieSDBackend unavailable ({e}); falling back to default torch.compile")
+            self.compile_backend = "default"
+            return None
 
     def get_compiled_block(self, block_idx, block):
         key = self.get_compile_block_key(block_idx, block)
@@ -21,7 +51,14 @@ def get_compiled_block(self, block_idx, block):
         def block_runner(*args):
             return self.infer_block(block, *args)
 
-        compiled = torch.compile(block_runner, dynamic=None)
+        compile_kwargs = {}
+        if self.compile_backend == "mindie" and self._compile_backend_obj is not None:
+            compile_kwargs["backend"] = self._compile_backend_obj
+        # dynamic=False: H3 per-step sequence is fixed (9467 local); fixed-shape
+        # graphs skip Dynamo's per-call shape guard checks. Fall back to
+        # dynamic=None if the graph contains truly dynamic inputs (Sym symbols).
+        dynamic_mode = getattr(self, "compile_dynamic", None)
+        compiled = torch.compile(block_runner, dynamic=dynamic_mode, **compile_kwargs)
         self.compiled_blocks[key] = (block, compiled)
         return compiled
 

lightx2v_platform/ops/rope/ascend_npu/__init__.py

@@ -1,3 +1,4 @@
+from .minimax_h3_npu_rope import MiniMaxH3NpuRope
 from .npu_rope import NpuRope
 
-__all__ = ["NpuRope"]
+__all__ = ["NpuRope", "MiniMaxH3NpuRope"]

lightx2v_platform/ops/rope/ascend_npu/minimax_h3_npu_rope.py

@@ -0,0 +1,75 @@
+"""Fused partial rotate-half RoPE for MiniMax-H3 on Ascend.
+
+MiniMax-H3 rotates only the leading ``rotary_dim`` (96 of 128) channels of each
+head with a rotate-half (split-half) pairing; the remaining channels pass
+through. The rotary part is fused into a single NPU kernel through the
+MindIE-SD ``rotary_position_embedding`` op (``npu_rotary_mul`` mode ``half``),
+and the pass-through channels are re-catenated outside.
+
+Input/output guards for the fused op live in MindIE-SD
+(``mindiesd.layers.rope.check_input_params``). Fallback: when mindiesd is not
+installed, delegate to ``TorchRealRope`` (split_half), the original decomposed
+flow.
+"""
+
+import torch
+
+from lightx2v_platform.ops.rope.rope_template import RopeTemplate
+from lightx2v_platform.registry_factory import PLATFORM_ROPE_REGISTER
+
+try:
+    from mindiesd.layers import rotary_position_embedding as _mindiesd_rope
+except ImportError:
+    _mindiesd_rope = None
+
+
+@PLATFORM_ROPE_REGISTER("minimax_h3_npu_rope")
+class MiniMaxH3NpuRope(RopeTemplate):
+    """Partial split-half RoPE fused via the MindIE-SD rotary_position_embedding op."""
+
+    def __init__(self, layout="split_half", compute_dtype=torch.float32):
+        super().__init__(layout=layout, compute_dtype=compute_dtype)
+        if layout != "split_half":
+            raise ValueError("MiniMaxH3NpuRope only supports split_half layout")
+
+    def _fallback(self):
+        from lightx2v.common.ops.rope import TorchRealRope
+
+        return TorchRealRope(layout=self.layout, compute_dtype=self.compute_dtype)
+
+    def apply(self, xq, xk, freqs, rotary_dim=None, **kwargs):
+        if _mindiesd_rope is None:
+            return self._fallback().apply(xq, xk, freqs, rotary_dim=rotary_dim, **kwargs)
+        return (
+            self._apply_single(xq, freqs, rotary_dim),
+            self._apply_single(xk, freqs, rotary_dim),
+        )
+
+    def apply_single(self, x, freqs, rotary_dim=None, **kwargs):
+        if _mindiesd_rope is None:
+            return self._fallback().apply_single(x, freqs, rotary_dim=rotary_dim, **kwargs)
+        return self._apply_single(x, freqs, rotary_dim)
+
+    def _apply_single(self, x, freqs, rotary_dim):
+        # x: [L, H, D]; rotate the leading rotary_dim channels, pass the rest.
+        cos, sin = freqs
+        rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim
+        x_rot = x[..., :rotary_dim].contiguous()
+        x_pass = x[..., rotary_dim:]
+        cos = cos.to(x.dtype).contiguous()
+        sin = sin.to(x.dtype).contiguous()
+        # mindiesd rotary_position_embedding takes x in [B,N,S,D]/[B,S,N,D]/
+        # [S,B,N,D] and 4-D cos/sin ([S,1,1,D] S11D is the SBND pairing);
+        # its 2-D [S,D] path assumes [B,S,N,D], so pass S11D explicitly.
+        rotated = _mindiesd_rope(
+            x_rot.unsqueeze(1),  # [L, 1, H, D] SBND
+            cos.unsqueeze(1).unsqueeze(1),  # [L, 1, 1, D] S11D
+            sin.unsqueeze(1).unsqueeze(1),
+            rotated_mode="rotated_half",
+            head_first=False,
+            fused=True,
+        ).squeeze(1)
+        rotated = rotated.to(x.dtype)
+        if x_pass.shape[-1]:
+            return torch.cat((rotated, x_pass), dim=-1)
+        return rotated