feat(mps): support MiniMax-H3 INT8 inference on Apple M5

configs/platforms/mps/minimax_h3_t2av_4step_int8.json

@@ -0,0 +1,61 @@
+{
+  "infer_steps": 4,
+  "num_frames": 124,
+  "size": [
+    544,
+    960
+  ],
+  "fps": 24,
+  "enable_cfg": false,
+  "cpu_offload": false,
+  "offload_granularity": "model",
+  "use_adaln_cache": false,
+  "text_encoder_cpu_offload": false,
+  "text_encoder_offload_granularity": "model",
+  "vae_cpu_offload": false,
+  "attn_type": "torch_sdpa_mps",
+  "rms_type": "torch_native",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": false,
+  "video_flow_shift": 6.0,
+  "audio_flow_shift": 3.0,
+  "h3_step_update": "training_euler",
+  "vae_spatial_scale_factor": 16,
+  "model_cls": "minimax_h3",
+  "model_variant": "fl2av",
+  "task": "t2av",
+  "local_files_only": true,
+  "mps_sdpa_query_chunk_size": 512,
+  "rope_compute_dtype": "bf16",
+  "dit_disk_streaming": false,
+  "text_encoder_disk_streaming": false,
+  "text_encoder_release_block_offload_buffers": false,
+  "warmup": false,
+  "lazy_load": false,
+  "shared_cpu_weights": false,
+  "tensor_parallel": false,
+  "text_encoder_tensor_parallel": false,
+  "seq_parallel": false,
+  "cfg_parallel": false,
+  "dit_release_block_offload_buffers": false,
+  "unload_modules": false,
+  "cpu_offload_block_contiguous": false,
+  "dit_quantized": true,
+  "dit_quant_scheme": "int8-metal",
+  "use_fused_qkv": false,
+  "text_encoder_shared_cpu_weights": false,
+  "text_encoder_quantized": false,
+  "text_encoder_load_on_init": true,
+  "video_vae_offload_granularity": "model",
+  "video_vae_shared_cpu_weights": false,
+  "video_vae_quantized": false,
+  "vae_use_compile": false,
+  "vae_encoder_conv_mode": "torch",
+  "vae_sensitive_layer_dtype": "fp32",
+  "vae_attn_type": "torch_sdpa",
+  "parallel": false,
+  "dit_quant_scope": "dit_proj",
+  "mps_resident_weights": true,
+  "text_encoder_load_vision_on_init": false
+}

lightx2v/common/offload/mps_resident.py

@@ -0,0 +1,62 @@
+"""Load resident weight containers directly into CPU-visible MPS storage."""
+
+import torch
+from loguru import logger
+
+from lightx2v.common.offload.mps_manager import host_view
+
+
+def weight_bindings(root):
+    seen, stack = set(), [root]
+    while stack:
+        obj = stack.pop()
+        if obj is None or id(obj) in seen:
+            continue
+        seen.add(id(obj))
+        stack.extend(getattr(obj, "_modules", {}).values())
+        stack.extend(getattr(obj, "_parameters", {}).values())
+        for name, attr, transpose in getattr(obj, "base_attrs", ()):
+            if attr == "weight":
+                transpose = getattr(obj, "weight_need_transpose", transpose)
+            yield obj, name, attr, transpose
+
+
+def load_resident_groups(checkpoint, groups):
+    """Adopt exact checkpoint dtypes, with no retained host weight copies.
+
+    Groups must partition the selected checkpoint. CPU aliases exist only while
+    reading idle allocations; synchronization orders both sides of host writes.
+    Container.load must adopt storage or a transpose view without conversion.
+    """
+    groups = list(groups)
+    names = [name for _, root in groups for _, name, _, _ in weight_bindings(root)]
+    if len(names) != len(set(names)) or set(names) != set(checkpoint.weight_map):
+        raise ValueError("Resident checkpoint and model bindings must match exactly")
+    total = 0
+    for label, root in groups:
+        bindings = list(weight_bindings(root))
+        tensors = {}
+        for _, name, _, _ in bindings:
+            dtype, shape, _, size = checkpoint.tensor_metadata(name)
+            tensors[name] = torch.empty(shape, dtype=dtype, device="mps")
+            total += size
+        torch.mps.synchronize()
+        aliases = {name: host_view(tensor) for name, tensor in tensors.items()}
+        checkpoint.load_tensors_into(aliases)
+        torch.mps.synchronize()
+        del aliases
+        root.load(tensors)
+        for obj, name, attr, transpose in bindings:
+            actual, original = getattr(obj, attr, None), tensors[name]
+            shape = tuple(reversed(original.shape)) if transpose else tuple(original.shape)
+            if (
+                not isinstance(actual, torch.Tensor)
+                or actual.device.type != "mps"
+                or actual.dtype != original.dtype
+                or tuple(actual.shape) != shape
+                or actual.untyped_storage()._cdata != original.untyped_storage()._cdata
+                or getattr(obj, "pin_" + attr, None) is not None
+            ):
+                raise RuntimeError("Resident loader requires exact storage adoption: " + name)
+        logger.info("MPS resident {} loaded; {:.3f} GiB cumulative weights", label, total / 1024**3)
+    return {"weight_bytes": total, "weight_count": len(names)}

lightx2v/common/ops/mm/optimized_linear.py

@@ -10,8 +10,6 @@
 import torch.nn.functional as F
 
 from lightx2v.common.ops.mm.fp8_f16_accum import cutlass_scaled_fp8_mm_f16_accum
-from lightx2v.common.ops.mm.fused_gelu_linear import fused_gelu_int8_convrot_linear, fused_gelu_quant_fp8_f16_accum, fused_gelu_quant_fp8_sgl, prepare_convrot_gelu_fusion
-from lightx2v.common.ops.mm.fused_swiglu_linear import fused_swiglu_quant_fp8_f16_accum
 from lightx2v.common.ops.mm.mm_weight import MMWeightTP, MMWeightWfp8channelAfp8channelF16Accum, MMWeightWfp8channelAfp8channelF32Accum, MMWeightWfp8channelAfp8channeldynamicSgl, MMWeightWint8ConvRot
 from lightx2v.common.ops.mm.projection_group import ConvRotProjectionGroup, FP8F16AccumProjectionGroup, FP8ProjectionGroup, _backend_apis_available
 from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm
@@ -79,6 +77,8 @@ def apply_swiglu_linear(module, value, gate):
         or (tp is not None and tp.split_dim != "row")
     ):
         return module.apply(value * F.silu(gate))
+    from lightx2v.common.ops.mm.fused_swiglu_linear import fused_swiglu_quant_fp8_f16_accum
+
     quantized, input_scale = fused_swiglu_quant_fp8_f16_accum(value, gate, qmax)
     output = linear.fp8_gemm(quantized, weight, input_scale, scale, linear.infer_dtype, linear._get_actual_bias())
     if tp is not None and tp.reduce_output and tp.tp_size > 1 and tp.tp_group is not None:
@@ -136,13 +136,19 @@ def apply(self, hidden):
         ):
             return module.apply(F.gelu(hidden, approximate="tanh"))
         if self.fusion == "fp8-sgl":
+            from lightx2v.common.ops.mm.fused_gelu_linear import fused_gelu_quant_fp8_sgl
+
             quantized, scale = fused_gelu_quant_fp8_sgl(hidden)
             return sgl_fp8_scaled_mm(quantized, self.weight, scale, self.scale, self.infer_dtype, self.bias)
         if self.fusion == "fp8-f16-accum":
+            from lightx2v.common.ops.mm.fused_gelu_linear import fused_gelu_quant_fp8_f16_accum
+
             quantized, scale = fused_gelu_quant_fp8_f16_accum(hidden, self.activation_qmax)
             return cutlass_scaled_fp8_mm_f16_accum(quantized, self.weight, scale, self.scale, self.infer_dtype, self.bias)
         if hidden.shape[0] < 3 or module._convrot_groupsize != 256:
             return module.apply(F.gelu(hidden, approximate="tanh"))
+        from lightx2v.common.ops.mm.fused_gelu_linear import fused_gelu_int8_convrot_linear
+
         return fused_gelu_int8_convrot_linear(hidden, self.weight, self.scale, self.bias, self.infer_dtype, module._convrot_groupsize)
 
 
@@ -175,6 +181,8 @@ def apply(self, hidden):
             or hidden.device != self.weight.device
         ):
             return module.apply(F.gelu(hidden, approximate="tanh"))
+        from lightx2v.common.ops.mm.fused_gelu_linear import fused_gelu_quant_fp8_f16_accum
+
         quantized, scale = fused_gelu_quant_fp8_f16_accum(hidden, self.activation_qmax)
         return cutlass_scaled_fp8_mm_f16_accum(quantized, self.weight, scale, self.scale, self.infer_dtype, self.bias)
 
@@ -233,6 +241,8 @@ def make_gelu_linear(module, *, enabled=True, offload_buffers=False):
         return GeluLinear(module)
     if type(module) is not MMWeightWint8ConvRot or module.infer_dtype != torch.bfloat16 or getattr(module, "_convrot_groupsize", None) != 256:
         return GeluLinear(module)
+    from lightx2v.common.ops.mm.fused_gelu_linear import prepare_convrot_gelu_fusion
+
     weight, scale, bias = getattr(module, "weight", None), getattr(module, "weight_scale", None), getattr(module, "bias", None)
     if (
         not isinstance(weight, torch.Tensor)

lightx2v/common/ops/mm/projection_group.py

@@ -10,7 +10,6 @@
 from lightx2v.common.ops.mm.fp8_f16_accum import cutlass_scaled_fp8_mm_f16_accum, fp8_f16_accum_mm_available
 from lightx2v.common.ops.mm.mm_weight import MMWeightWfp8channelAfp8channelF16Accum, MMWeightWfp8channelAfp8channeldynamicSgl, MMWeightWint8ConvRot
 from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_per_token_quant, sgl_fp8_scaled_mm, sgl_kernel
-from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton
 
 try:
     import comfy_kitchen
@@ -342,6 +341,8 @@ def __init__(self, modules, packed=False, offload_buffers=False):
         self.activation_qmax = modules[0].fp8_activation_qmax
 
     def quantize(self, x):
+        from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton
+
         return fp8_quantize_range_triton(x, self.activation_qmax)
 
     def matmul(self, quantized, weight, input_scale, weight_scale, bias):

lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py

@@ -1021,6 +1021,26 @@ def load_text_encoder(self):
             block_offload=self.block_offload,
             tp_group=self.tp_group,
         )
+        if self.config.get("mps_resident_weights", False):
+            if AI_DEVICE != "mps" or self.cpu_offload or self.tensor_parallel or quantized or self.config.get("text_encoder_shared_cpu_weights", False):
+                raise ValueError("MPS resident H3 text loading requires unquantized single-device resident text weights")
+            from lightx2v.common.offload.mps_resident import load_resident_groups
+            from lightx2v.common.offload.safetensors_checkpoint import SafetensorsCheckpoint
+
+            mapping, dtype = self._preflight_native_checkpoint(text_encoder, text_encoder_path, text_config)
+            if dtype != "BF16":
+                raise ValueError("MPS resident H3 text loading requires BF16 checkpoint weights")
+            embedding = text_encoder.embed_tokens
+            embedding.base_attrs = [(embedding.weight_name, "weight", False)]
+            groups = [("text embedding", embedding)] + [(f"text layer {i}", layer) for i, layer in enumerate(text_encoder.layers)]
+            text_encoder._inv_freq = text_encoder._inv_freq.to("mps")
+            # Indexed releases also contain the unused decoder suffix and
+            # vision tower. Only the validated native 50-layer prefix is resident.
+            selected = {name: mapping[name] for name, _ in text_encoder.named_weight_modules()}
+            checkpoint = SafetensorsCheckpoint(text_encoder_path, weight_map=selected)
+            self.mps_resident_storage = load_resident_groups(checkpoint, groups)
+            self.text_encoder = text_encoder
+            return text_encoder
         if self.config.get("text_encoder_shared_cpu_weights", False):
             from lightx2v.models.input_encoders.hf.minimax_h3.shared_weights import load_shared_text_weights
 
@@ -1082,9 +1102,11 @@ def to_cpu(self):
 
     def load(self):
         self.load_tokenizer()
-        self.load_processor()
+        if self.config.get("text_encoder_load_vision_on_init", True):
+            self.load_processor()
         self.load_text_encoder()
-        self.load_vision_encoder()
+        if self.config.get("text_encoder_load_vision_on_init", True):
+            self.load_vision_encoder()
         return self
 
     def unload(self):

lightx2v/models/networks/minimax_h3/int8_metal.py

@@ -0,0 +1,104 @@
+"""Mixed-precision selection and checkpoint contract for H3 Apple INT8."""
+
+import json
+from pathlib import Path
+
+MLP_SUFFIXES = (".ff.net.0.proj", ".ff.net.2")
+QUANT_SCOPES = {
+    "fc1": ("fc1",),
+    "fc2": ("fc2",),
+    "mlp": ("fc1", "fc2"),
+    "mlp_out": ("fc1", "fc2", "out"),
+    "mlp_qkv": ("fc1", "fc2", "qkv"),
+    "dit_proj": ("fc1", "fc2", "qkv", "out"),
+}
+
+
+def linear_component(name):
+    if not name.startswith("transformer_blocks."):
+        return None
+    if name.endswith(".ff.net.0.proj"):
+        return "fc1"
+    if name.endswith(".ff.net.2"):
+        return "fc2"
+    if name.endswith((".attn.to_q", ".attn.to_k", ".attn.to_v")):
+        return "qkv"
+    if name.endswith(".attn.to_out.0"):
+        return "out"
+    return None
+
+
+def is_mlp_linear(name):
+    return name.startswith("transformer_blocks.") and name.endswith(MLP_SUFFIXES)
+
+
+def linear_quant_scheme(config, name):
+    scheme = config.get("dit_quant_scheme", "Default")
+    scope = config.get("dit_quant_scope", "all")
+    if scope == "all":
+        return scheme
+    if scope not in QUANT_SCOPES or scheme != "int8-metal":
+        raise ValueError("H3 dit_quant_scope supports all, or an explicit int8-metal projection profile")
+    return scheme if linear_component(name) in QUANT_SCOPES[scope] else "Default"
+
+
+def validate_metal_checkpoint(config):
+    path = Path(config["dit_quantized_ckpt"])
+    indices = list(path.glob("*.safetensors.index.json")) if path.is_dir() else []
+    if len(indices) != 1:
+        raise ValueError("H3 int8-metal requires one indexed safetensors checkpoint directory")
+    index = json.loads(indices[0].read_text())
+    metadata = index.get("metadata", {})
+    if metadata.get("lightx2v_quant_scheme") != "int8-metal":
+        raise ValueError("H3 checkpoint does not declare lightx2v_quant_scheme=int8-metal")
+    if metadata.get("lightx2v_quant_scope") != config.get("dit_quant_scope", "all"):
+        raise ValueError("H3 INT8 checkpoint scope does not match dit_quant_scope")
+    import torch
+
+    from lightx2v.common.offload.safetensors_checkpoint import SafetensorsCheckpoint
+
+    checkpoint = SafetensorsCheckpoint(path, weight_map=index["weight_map"])
+    scope = config.get("dit_quant_scope", "all")
+    selected = set()
+    for name in checkpoint.weight_map:
+        dtype, shape, _, _ = checkpoint.tensor_metadata(name)
+        if name.endswith(".weight") and linear_quant_scheme(dict(config, dit_quant_scheme="int8-metal"), name[:-7]) == "int8-metal":
+            # Explicit scopes select only their declared projections.
+            if scope == "all" and linear_component(name[:-7]) is None:
+                continue
+            if dtype != torch.int8 or len(shape) != 2 or any(n <= 0 or n % 128 for n in shape):
+                raise ValueError("Invalid INT8 Metal weight shape/dtype: " + name)
+            scale = name + "_scale"
+            if scale not in checkpoint.weight_map:
+                raise ValueError("Missing INT8 Metal scale: " + scale)
+            sdtype, sshape, _, _ = checkpoint.tensor_metadata(scale)
+            if sdtype != torch.float32 or sshape != (shape[0], 1):
+                raise ValueError("Invalid INT8 Metal scale shape/dtype: " + scale)
+            selected.add(name)
+        elif dtype == torch.int8:
+            raise ValueError("INT8 tensor outside declared projection scope: " + name)
+    if not selected:
+        raise ValueError("No INT8 weights in declared H3 projection scope")
+    if metadata.get("quantized_linear_count", len(selected)) != len(selected):
+        raise ValueError("Declared INT8 projection count does not match checkpoint")
+    return index
+
+
+def load_metal_checkpoint(config, device, remove_keys=()):
+    """Honor the authoritative index, including selective checkpoint views."""
+    from collections import defaultdict
+
+    from safetensors import safe_open
+
+    index = validate_metal_checkpoint(config)
+    root = Path(config["dit_quantized_ckpt"])
+    by_file = defaultdict(list)
+    for name, file in index["weight_map"].items():
+        if not any(part in name for part in remove_keys):
+            by_file[file].append(name)
+    tensors = {}
+    for file, names in sorted(by_file.items()):
+        with safe_open(str(root / file), framework="pt", device="cpu") as source:
+            for name in names:
+                tensors[name] = source.get_tensor(name).to(device)
+    return tensors