Optimize MiniMax-H3

configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json

@@ -0,0 +1,43 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 362,
+  "target_height": 768,
+  "target_width": 1344,
+  "fps": 24,
+  "target_fps": 24,
+  "enable_cfg": false,
+  "cpu_offload": true,
+  "offload_granularity": "block",
+  "text_encoder_cpu_offload": true,
+  "text_encoder_offload_granularity": "block",
+  "vae_cpu_offload": true,
+  "lazy_load": false,
+  "unload_modules": false,
+  "attn_type": "dynamic_sparse_attn",
+  "dynamic_sparse_attn_setting": {
+      "sparsity_ratio": 0.9,
+      "operator": "sage2"
+  },
+  "rms_type": "sgl-kernel",
+  "rope_type": "minimax_h3_triton_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "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,
+  "dit_quantized": true,
+  "dit_quant_scheme": "fp8-sgl",
+  "dit_quantized_ckpt": "/path/to/models/minimax_h3/h3_quantized/fp8/minimax_h3_fp8.safetensors",
+  "parallel": {
+    "seq_p_size": 8,
+    "seq_p_attn_type": "ulysses",
+    "seq_p_prepost_backend": "triton",
+    "seq_p_tensor_fusion": true,
+    "seq_p_head_parallel": true,
+    "seq_p_fp8_comm": true
+  }
+}

lightx2v/common/ops/attn/dynamic_sparse_attn.py

@@ -36,8 +36,14 @@ class DynamicSparseAttnWeight(AttnWeightTemplate):
     operator = "triton"
     per_block_mean = False
 
-    def __init__(self):
-        self.config = {}
+    def __init__(self, config=None):
+        self.config = dict(config or {})
+        self.sparsity_ratio = float(self.config.get("sparsity_ratio", type(self).sparsity_ratio))
+        self.operator = self.config.get("operator", type(self).operator)
+        self.per_block_mean = bool(self.config.get("per_block_mean", type(self).per_block_mean))
+
+        if not 0.0 <= self.sparsity_ratio < 1.0:
+            raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}")
 
         self.arch = get_cuda_arch(torch.cuda.current_device())
         self.topk = 1 - self.sparsity_ratio
@@ -65,7 +71,7 @@ def __init__(self):
         else:
             raise NotImplementedError(f"Not supported SLA operator: {self.operator}.")
 
-        logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}")
+        # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}")
 
     def apply(
         self,

lightx2v/models/networks/minimax_h3/weights/pre_weights.py

@@ -39,7 +39,13 @@ def __init__(self, prefix, config):
             "norm_k",
             _rms(config, f"{prefix}.norm_k.weight", eps=float(config.get("qk_norm_eps", 1e-5))),
         )
-        self.add_module("calculate", ATTN_WEIGHT_REGISTER[config.get("attn_type", "flash_attn3")]())
+        attn_type = config.get("attn_type", "flash_attn3")
+        attention_cls = ATTN_WEIGHT_REGISTER[attn_type]
+        if attn_type == "dynamic_sparse_attn":
+            calculate = attention_cls(config.get("dynamic_sparse_attn_setting", {}))
+        else:
+            calculate = attention_cls()
+        self.add_module("calculate", calculate)
         self.add_module("to_out", _linear(f"{prefix}.to_out.0", config=config, tp_split="row"))
 
 

lightx2v/models/networks/minimax_h3/weights/transformer_weights.py

@@ -67,7 +67,13 @@ def __init__(self, prefix, config, create_cuda_buffer=False):
                 compute_dtype=torch.float32,
             ),
         )
-        self.add_module("calculate", ATTN_WEIGHT_REGISTER[config.get("attn_type", "flash_attn3")]())
+        attn_type = config.get("attn_type", "flash_attn3")
+        attention_cls = ATTN_WEIGHT_REGISTER[attn_type]
+        if attn_type == "dynamic_sparse_attn":
+            calculate = attention_cls(config.get("dynamic_sparse_attn_setting", {}))
+        else:
+            calculate = attention_cls()
+        self.add_module("calculate", calculate)
         if config.get("seq_parallel", False):
             parallel = config.get("parallel", {})
             self.add_module(

scripts/minimax_h3/run_minimax_h3_t2av_sp.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=/path/to/LightX2V
+model_path=/path/to/models/minimax_h3/h3_hf_bf16
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+
+# set environment variables
+source ${lightx2v_path}/scripts/base/base.sh
+export DTYPE=BF16
+export SENSITIVE_LAYER_DTYPE=FP32
+
+torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \
+--model_cls minimax_h3 \
+--task t2av \
+--model_path $model_path \
+--config_json ${lightx2v_path}/configs/minimax_h3/fp8/minimax_h3_t2av_sp8_5090.json \
+--prompt "A cinematic fox walking through a snowy forest" \
+--save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_sp.mp4 \
+--seed 42