diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json new file mode 100644 index 000000000..77736c6a8 --- /dev/null +++ b/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"}, +} diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json new file mode 100644 index 000000000..c73cba50d --- /dev/null +++ b/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"}, +} diff --git a/lightx2v/common/ops/attn/ulysses_a2a.py b/lightx2v/common/ops/attn/ulysses_a2a.py index 22cec6fdf..4bdf99a04 100644 --- a/lightx2v/common/ops/attn/ulysses_a2a.py +++ b/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 diff --git a/lightx2v/common/transformer_infer/transformer_infer.py b/lightx2v/common/transformer_infer/transformer_infer.py index cff6b7d44..28dd3d1e9 100644 --- a/lightx2v/common/transformer_infer/transformer_infer.py +++ b/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 diff --git a/lightx2v_platform/ops/rope/ascend_npu/__init__.py b/lightx2v_platform/ops/rope/ascend_npu/__init__.py index ba664c3de..79e04d84d 100644 --- a/lightx2v_platform/ops/rope/ascend_npu/__init__.py +++ b/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"] diff --git a/lightx2v_platform/ops/rope/ascend_npu/minimax_h3_npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/minimax_h3_npu_rope.py new file mode 100644 index 000000000..d62c2aec0 --- /dev/null +++ b/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