feat(minimax_h3): add MPS support, Turbo config and VAE fixes

configs/platforms/mps/minimax_h3_t2av.json

@@ -0,0 +1,51 @@
+{
+  "infer_steps": 29,
+  "num_frames": 124,
+  "size": [480, 480],
+  "fps": 24,
+  "enable_cfg": false,
+
+  "cpu_offload": true,
+  "offload_granularity": "block",
+  "dit_prepost_resident": false,
+  "dit_disk_streaming": true,
+  "dit_mps_shared_buffer": true,
+  "use_adaln_cache": true,
+  "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers",
+
+  "text_encoder_cpu_offload": true,
+  "text_encoder_offload_granularity": "block",
+  "text_encoder_disk_streaming": true,
+  "text_encoder_host_pinned": false,
+  "text_encoder_release_block_offload_buffers": true,
+  "text_encoder_quantized": false,
+
+  "vae_cpu_offload": true,
+  "vae_use_compile": false,
+  "vae_attn_type": "torch_sdpa",
+  "video_vae_quantized": false,
+
+  "lazy_load": false,
+  "unload_modules": false,
+  "warmup": false,
+
+  "attn_type": "torch_sdpa",
+  "mps_sdpa_query_chunk_size": 512,
+  "rms_type": "torch_native",
+  "rope_type": "torch_real_rope",
+
+  "tensor_parallel": false,
+  "dit_quantized": false,
+  "dit_quant_scheme": "Default",
+
+  "feature_caching": "NoCaching",
+  "use_compile": false,
+
+  "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
+}

configs/platforms/mps/minimax_h3_t2av_4step_512_22.json

@@ -0,0 +1,51 @@
+{
+  "infer_steps": 4,
+  "num_frames": 22,
+  "size": [512, 512],
+  "fps": 24,
+  "enable_cfg": false,
+
+  "cpu_offload": true,
+  "offload_granularity": "block",
+  "dit_prepost_resident": false,
+  "dit_disk_streaming": true,
+  "dit_mps_shared_buffer": true,
+  "use_adaln_cache": true,
+  "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers",
+
+  "text_encoder_cpu_offload": true,
+  "text_encoder_offload_granularity": "block",
+  "text_encoder_disk_streaming": true,
+  "text_encoder_host_pinned": false,
+  "text_encoder_release_block_offload_buffers": true,
+  "text_encoder_quantized": false,
+
+  "vae_cpu_offload": true,
+  "vae_use_compile": false,
+  "vae_attn_type": "torch_sdpa",
+  "video_vae_quantized": false,
+
+  "lazy_load": false,
+  "unload_modules": false,
+  "warmup": false,
+
+  "attn_type": "torch_sdpa",
+  "mps_sdpa_query_chunk_size": 512,
+  "rms_type": "torch_native",
+  "rope_type": "torch_real_rope",
+
+  "tensor_parallel": false,
+  "dit_quantized": false,
+  "dit_quant_scheme": "Default",
+
+  "feature_caching": "NoCaching",
+  "use_compile": false,
+
+  "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
+}

lightx2v/common/offload/manager.py

@@ -1,4 +1,5 @@
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext
 
 import torch
 from loguru import logger
@@ -14,9 +15,12 @@
 class WeightAsyncStreamManager(object):
     def __init__(self, offload_granularity):
         self.offload_granularity = offload_granularity
-        self.init_stream = torch_device_module.Stream(priority=0)
         self.need_init_first_buffer = True
         self.lazy_load = False
+        self._init_streams()
+
+    def _init_streams(self):
+        self.init_stream = torch_device_module.Stream(priority=0)
         torch_version = parse(torch.__version__.split("+")[0])
         # Legacy name: this is the active device backend's weight-loading stream, not a CUDA-only stream.
         if AI_DEVICE == "cuda" and torch_version >= parse("2.7"):
@@ -26,6 +30,14 @@ def __init__(self, offload_granularity):
             self.cuda_load_stream = torch_device_module.Stream(priority=0)
             self.compute_stream = torch_device_module.Stream(priority=-1)
 
+    def prepare_compute(self):
+        self.compute_stream.wait_stream(torch_device_module.current_stream())
+
+    def compute_context(self):
+        if AI_DEVICE == "xpu":
+            return nullcontext()
+        return torch_device_module.stream(self.compute_stream)
+
     def init_cpu_buffer(self, blocks_cpu_buffer=None, phases_cpu_buffer=None):
         self.need_init_first_buffer = True
         if self.offload_granularity == "block":

lightx2v/common/offload/mps_manager.py

@@ -0,0 +1,89 @@
+"""Double-buffered disk offload into CPU-visible MPS weight storage."""
+
+from contextlib import nullcontext
+
+import torch
+from loguru import logger
+
+from lightx2v.common.offload.manager import WeightAsyncStreamManager
+
+
+def host_view(tensor):
+    """Alias MPS storage without copying; the caller orders CPU/GPU accesses."""
+    if tensor.device.type != "mps":
+        raise ValueError("Shared weight buffers require an MPS tensor")
+    if not hasattr(torch.mps, "_host_alias_storage"):
+        raise RuntimeError("Shared MPS weight buffers require torch.mps._host_alias_storage (PyTorch >= 2.13)")
+    storage = torch.mps._host_alias_storage(tensor.untyped_storage())
+    return torch.empty(0, dtype=tensor.dtype).set_(storage, tensor.storage_offset(), tensor.shape, tensor.stride())
+
+
+class MpsSharedWeightAsyncStreamManager(WeightAsyncStreamManager):
+    """Read into the idle slot while the default MPS stream uses the other.
+
+    The source implements load_block_into(target_buffer, block_idx). This runs
+    on a CPU worker and must only write host views of the supplied idle buffer.
+    """
+
+    def __init__(self, offload_granularity="block"):
+        if offload_granularity != "block":
+            raise ValueError("Shared MPS weight offload only supports block granularity")
+        super().__init__(offload_granularity)
+        self.cuda_buffers = []
+        self.init_lazy_load(num_workers=1)
+        logger.info("MPS shared weight offload: two device buffers, direct file reads, one prefetch worker")
+
+    def _init_streams(self):
+        # Disk I/O runs on a CPU thread; GPU work stays on the default stream.
+        pass
+
+    def prepare_compute(self):
+        pass
+
+    def compute_context(self):
+        return nullcontext()
+
+    def init_cuda_buffer(self, blocks_cuda_buffer=None, phases_cuda_buffer=None):
+        if blocks_cuda_buffer is None or len(blocks_cuda_buffer) != 2:
+            raise ValueError("Shared MPS offload requires exactly two device buffers")
+        if self.prefetch_futures:
+            raise RuntimeError("Cannot replace shared buffers while a prefetch is pending")
+        super().init_cuda_buffer(blocks_cuda_buffer, phases_cuda_buffer)
+        if self.executor is None:
+            self.init_lazy_load(num_workers=1)
+
+    def init_first_buffer(self, blocks, adapter_block_idx=None):
+        torch.mps.synchronize()
+        blocks.load_block_into(self.cuda_buffers[0], 0)
+        torch.mps.synchronize()
+        self.need_init_first_buffer = False
+
+    def prefetch_weights(self, block_idx, blocks, adapter_block_idx=None):
+        if self.prefetch_futures:
+            raise RuntimeError("Call swap_blocks before scheduling another prefetch")
+        self.prefetch_block_idx = block_idx
+        self.prefetch_futures = [self.executor.submit(blocks.load_block_into, self.cuda_buffers[1], block_idx)]
+
+    def swap_blocks(self):
+        if not self.prefetch_futures:
+            raise RuntimeError("No shared-buffer prefetch to complete")
+        # This commits GPU work and releases the GIL while waiting, so the disk
+        # reader continues filling slot 1 while the GPU consumes slot 0.
+        torch.mps.synchronize()
+        for future in self.prefetch_futures:
+            future.result()
+        torch.mps.synchronize()
+        self.prefetch_futures.clear()
+        self.cuda_buffers.reverse()
+
+    def close(self):
+        """Drain both users before releasing aliases; allow later reinitialization."""
+        try:
+            if self.executor is not None:
+                self.executor.shutdown(wait=True)
+            torch.mps.synchronize()
+        finally:
+            self.executor = None
+            self.prefetch_futures.clear()
+            self.cuda_buffers = []
+            self.need_init_first_buffer = True

lightx2v/common/ops/attn/__init__.py

@@ -1,22 +1,26 @@
-from .draft_attn import DraftAttnWeight
-from .dynamic_sparse_attn import DynamicSparseAttnWeight
-from .flash_attn import (
-    FlashAttn2Weight,
-    FlashAttn3Weight,
-    FlashAttn4Weight,
-    SparseFlashAttn4Weight,
-)
-from .general_sparse_attn import GeneralSparseAttnWeight
-from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer
-from .radial_attn import RadialAttnWeight
-from .rainfusion_attn import RainfusionAttnWeight
-from .ring_attn import RingAttnWeight
-from .sage_attn import SageAttn2KInt8VFP8Weight, SageAttn2Weight, SageAttn3Weight, SparseSageAttn2Weight, SparseSageAttn3Weight
-from .sol_attn import SolAttnWeight
-from .sparge_attn import SpargeAttnWeight
-from .sparse_mask_generator import NbhdMaskGenerator, SlaMaskGenerator, SpargeMaskGenerator, SvgMaskGenerator
-from .sparse_operator import FlashinferOperator, FlexBlockOperator, MagiOperator, SlaTritonOperator, SparseFlashAttentionV4Operator, SparseSageAttentionV2Operator, SparseSageAttentionV3Operator
-from .svg2_attn import Svg2AttnWeight
-from .svg_attn import SvgAttnWeight
+from lightx2v_platform.base.global_var import AI_DEVICE
+
 from .torch_sdpa import TorchSDPAWeight
-from .ulysses_attn import UlyssesAttnWeight
+
+if str(AI_DEVICE) != "mps":
+    from .draft_attn import DraftAttnWeight
+    from .dynamic_sparse_attn import DynamicSparseAttnWeight
+    from .flash_attn import (
+        FlashAttn2Weight,
+        FlashAttn3Weight,
+        FlashAttn4Weight,
+        SparseFlashAttn4Weight,
+    )
+    from .general_sparse_attn import GeneralSparseAttnWeight
+    from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer
+    from .radial_attn import RadialAttnWeight
+    from .rainfusion_attn import RainfusionAttnWeight
+    from .ring_attn import RingAttnWeight
+    from .sage_attn import SageAttn2KInt8VFP8Weight, SageAttn2Weight, SageAttn3Weight, SparseSageAttn2Weight, SparseSageAttn3Weight
+    from .sol_attn import SolAttnWeight
+    from .sparge_attn import SpargeAttnWeight
+    from .sparse_mask_generator import NbhdMaskGenerator, SlaMaskGenerator, SpargeMaskGenerator, SvgMaskGenerator
+    from .sparse_operator import FlashinferOperator, FlexBlockOperator, MagiOperator, SlaTritonOperator, SparseFlashAttentionV4Operator, SparseSageAttentionV2Operator, SparseSageAttentionV3Operator
+    from .svg2_attn import Svg2AttnWeight
+    from .svg_attn import SvgAttnWeight
+    from .ulysses_attn import UlyssesAttnWeight

lightx2v/common/ops/attn/torch_sdpa.py

@@ -8,6 +8,33 @@
 from .template import AttnWeightTemplate
 
 
+def _use_h3_mps_query_chunks(q, k, v, chunk_size, scope, attn_mask, causal, drop_rate):
+    return (
+        scope == "minimax_h3_dit"
+        and isinstance(chunk_size, int)
+        and chunk_size > 0
+        and q.device.type == k.device.type == v.device.type == "mps"
+        and q.ndim == 4
+        and q.shape == k.shape == v.shape
+        and q.shape[0] == 1
+        and q.shape[1] == 56
+        and q.shape[-1] == 128
+        and q.shape[2] > 0
+        and attn_mask is None
+        and not causal
+        and drop_rate == 0
+    )
+
+
+def _query_chunked_sdpa(q, k, v, chunk_size):
+    # Each query still attends to every key/value. Only the query workspace is
+    # bounded; there is no context truncation or change to the softmax domain.
+    return torch.cat(
+        [F.scaled_dot_product_attention(q[:, :, start : start + chunk_size, :], k, v, attn_mask=None, dropout_p=0.0, is_causal=False) for start in range(0, q.shape[2], chunk_size)],
+        dim=2,
+    )
+
+
 @ATTN_WEIGHT_REGISTER("torch_sdpa")
 class TorchSDPAWeight(AttnWeightTemplate):
     def __init__(self):
@@ -44,17 +71,21 @@ def apply(
                 enable_mem_efficient=True,
             )
         with sdpa_ctx:
-            # q/k/v are (B, H, S, D) here, so head count is dim 1. GQA models such as
-            # neopp (32 q heads, 8 kv heads) need SDPA to broadcast the kv groups.
-            x = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                attn_mask=attn_mask,
-                dropout_p=drop_rate,
-                is_causal=causal,
-                enable_gqa=q.shape[1] != k.shape[1],
-            )
+            chunk_size = kwargs.get("mps_sdpa_query_chunk_size", 0)
+            if _use_h3_mps_query_chunks(q, k, v, chunk_size, kwargs.get("attention_scope"), attn_mask, causal, drop_rate):
+                x = _query_chunked_sdpa(q, k, v, chunk_size)
+            else:
+                # q/k/v are (B, H, S, D) here, so head count is dim 1. GQA models such as
+                # neopp (32 q heads, 8 kv heads) need SDPA to broadcast the kv groups.
+                x = F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=attn_mask,
+                    dropout_p=drop_rate,
+                    is_causal=causal,
+                    enable_gqa=q.shape[1] != k.shape[1],
+                )
         x = x.transpose(1, 2)
         b, s, a, d = x.shape
         out = x.reshape(b, s, -1)

lightx2v/common/ops/mm/fp8_f16_accum.py

@@ -2,8 +2,6 @@
 
 import torch
 
-from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton
-
 try:
     from lightx2v_kernel.gemm import FP8_F16_ACCUM_MM_AVAILABLE, cutlass_scaled_fp8_mm_f16_accum
 except ImportError:
@@ -35,6 +33,8 @@ def validate_fp8_f16_accum_qmax(activation_qmax):
 
 
 def fp8_f16_accum_linear(input_tensor, weight, weight_scale, bias, activation_qmax):
+    from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton
+
     input_shape = input_tensor.shape
     input_matrix = input_tensor.reshape(-1, input_shape[-1])
     quantized, activation_scale = fp8_quantize_range_triton(input_matrix, activation_qmax)