[diffusion] optimization: reduce minimax h3 mps memory pressure (#33880)

python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py

@@ -25,6 +25,9 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
     AttentionRequirements,
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
+    LAYERWISE_OFFLOAD,
+)
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
@@ -186,6 +189,29 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
     def validate_server_args(self, server_args) -> None:
         # Reject known-inexact VAE modes before any large component download.
         self.vae_config.resolved_parallel_decode_mode()
+        if current_platform.is_mps():
+            required_components = (
+                "transformer",
+                "text_encoder",
+                "video_vae",
+                "audio_vae",
+            )
+            missing_components = [
+                component
+                for component in required_components
+                if server_args.residency_mode(component) != LAYERWISE_OFFLOAD
+            ]
+            if missing_components:
+                raise ValueError(
+                    "MiniMax-H3 on MPS requires synchronous layerwise offload for "
+                    f"{missing_components}; pass --layerwise-offload-components "
+                    "transformer text_encoder video_vae audio_vae"
+                )
+            if server_args.enable_torch_compile:
+                raise ValueError(
+                    "MiniMax-H3 MPS execution does not support torch.compile; "
+                    "pass --enable-torch-compile false"
+                )
         component_backends = server_args.component_attention_backends or {}
         attention_backend = component_backends.get(
             "transformer", self._server_arg_value(server_args.attention_backend)

python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py

@@ -24,6 +24,8 @@ _PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
     SDPBackend.MATH,
 ]
 
+_MPS_VARLEN_QUERY_CHUNK_SIZE = 128
+
 
 class SDPABackend(AttentionBackend):
 
@@ -128,13 +130,31 @@ class SDPAImpl(AttentionImpl):
         for start, stop in zip(bounds[:-1], bounds[1:]):
             if start == stop:
                 continue
-            segment = self.forward(
+            if query.device.type != "mps":
-                query[start:stop].unsqueeze(0),
+                segment = self.forward(
-                key[start:stop].unsqueeze(0),
+                    query[start:stop].unsqueeze(0),
-                value[start:stop].unsqueeze(0),
+                    key[start:stop].unsqueeze(0),
-                None,
+                    value[start:stop].unsqueeze(0),
-            )
+                    None,
-            output[start:stop].copy_(segment[0])
+                )
+                output[start:stop].copy_(segment[0])
+                continue
+
+            # mps SDPA materializes a quadratic temporary for a varlen segment
+            # chunking query rows keeps every row's complete K/V context intact
+            keys = key[start:stop].unsqueeze(0)
+            values = value[start:stop].unsqueeze(0)
+            for query_start in range(start, stop, _MPS_VARLEN_QUERY_CHUNK_SIZE):
+                query_stop = min(query_start + _MPS_VARLEN_QUERY_CHUNK_SIZE, stop)
+                segment = self.forward(
+                    query[query_start:query_stop].unsqueeze(0),
+                    keys,
+                    values,
+                    None,
+                )
+                output[query_start:query_stop].copy_(segment[0])
+                torch.mps.synchronize()
+                torch.mps.empty_cache()
         return output
 
 

python/sglang/multimodal_gen/runtime/layers/linear.py

@@ -154,11 +154,11 @@ class UnquantizedLinearMethod(LinearMethodBase):
     def apply(
         self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
     ) -> torch.Tensor:
-        if x.device.type == "mps" and (
+        if x.device.type == "mps":
-            x.dtype != torch.float32
+            if x.dtype == layer.weight.dtype and (
-            or layer.weight.dtype != torch.float32
+                bias is None or bias.dtype == x.dtype
-            or (bias is not None and bias.dtype != torch.float32)
+            ):
-        ):
+                return F.linear(x, layer.weight, bias)
             return F.linear(
                 x.to(torch.float32),
                 layer.weight.to(torch.float32),

python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py

@@ -530,7 +530,7 @@ class TextEncoderLoader(ComponentLoader):
             )
             component_starts_on_cpu = False
 
-        if component_starts_on_cpu and not current_platform.is_mps():
+        if component_starts_on_cpu:
             model_device = torch.device("cpu")
         else:
             model_device = local_torch_device
@@ -559,6 +559,12 @@ class TextEncoderLoader(ComponentLoader):
                 )
             model.bind_encoder_tp_group(encoder_tp_group)
 
+            if current_platform.is_mps() and component_starts_on_cpu:
+                # the h3 encoder is layered immediately after this loader returns
+                # compatible CPU safetensors stay mapped instead of copying the
+                # full Qwen checkpoint into unified memory
+                model._mps_zero_copy_weight_loading = True
+
             weights_to_load = {name for name, _ in model.named_parameters()}
             loaded_weights = model.load_weights(
                 self._get_all_weights(
@@ -581,7 +587,10 @@ class TextEncoderLoader(ComponentLoader):
 
             if component_starts_on_cpu:
                 if current_platform.is_mps():
-                    model = model.to(local_torch_device)
+                    logger.info(
+                        "Keeping %s on CPU for MPS layerwise offload",
+                        model.__class__.__name__,
+                    )
                 else:
                     model = model.to("cpu")
             else:

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -136,6 +136,18 @@ class TransformerLoader(ComponentLoader):
     ]
     expected_library = "diffusers"
 
+    def customized_load_kwargs_for_component(
+        self, server_args: ServerArgs, component_name: str
+    ) -> dict[str, bool]:
+        if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
+            server_args, component_name
+        ):
+            logger.info(
+                "Loading %s on CPU first for MPS layerwise offload", component_name
+            )
+            return {"cpu_offload_flag": True}
+        return {}
+
     def should_raise_customized_load_error(
         self, server_args: ServerArgs, component_name: str
     ) -> bool:
@@ -151,7 +163,11 @@ class TransformerLoader(ComponentLoader):
         )
 
     def load_customized(
-        self, component_model_path: str, server_args: ServerArgs, component_name: str
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        cpu_offload_flag: bool = False,
     ):
         """Load the transformer based on the model path, and inference args."""
         component_server_args = _server_args_for_transformer_component(
@@ -196,8 +212,9 @@ class TransformerLoader(ComponentLoader):
         )
         # Quantization adapters may require resident weights, so placement must
         # be resolved after they have validated the component configuration.
-        component_starts_on_cpu = server_args.should_start_component_on_cpu(
+        component_starts_on_cpu = (
-            component_name
+            server_args.should_start_component_on_cpu(component_name)
+            or cpu_offload_flag
         )
         use_fsdp = server_args.should_use_fsdp_for_component(component_name)
 
@@ -259,10 +276,14 @@ class TransformerLoader(ComponentLoader):
             logger.debug("quantization config: %s", init_params["quant_config"])
 
         local_torch_device = get_local_torch_device()
-        checkpoint_load_device = _resolve_checkpoint_load_device(
+        checkpoint_load_device = (
-            local_torch_device,
+            torch.device("cpu")
-            component_starts_on_cpu=component_starts_on_cpu,
+            if cpu_offload_flag
-            runtime_quant_config=quant_spec.runtime_quant_config,
+            else _resolve_checkpoint_load_device(
+                local_torch_device,
+                component_starts_on_cpu=component_starts_on_cpu,
+                runtime_quant_config=quant_spec.runtime_quant_config,
+            )
         )
         direct_gpu_weight_loading = bool(
             component_server_args.direct_gpu_weight_loading
@@ -276,6 +297,9 @@ class TransformerLoader(ComponentLoader):
             needs_device_weight_postprocess=quant_spec.needs_device_weight_postprocess,
             component_starts_on_cpu=component_starts_on_cpu,
             load_full_state_dict_on_device=direct_gpu_weight_loading,
+            mps_layerwise_cpu_staging=bool(
+                cpu_offload_flag and current_platform.is_mps()
+            ),
         )
         if direct_gpu_weight_loading:
             logger.warning(

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -98,8 +98,24 @@ class VAELoader(ComponentLoader):
     component_names = ["vae", "audio_vae", "video_vae"]
     expected_library = "diffusers"
 
+    def customized_load_kwargs_for_component(
+        self, server_args: ServerArgs, component_name: str
+    ) -> dict[str, bool]:
+        if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
+            server_args, component_name
+        ):
+            logger.info(
+                "Loading %s on CPU first for MPS layerwise offload", component_name
+            )
+            return {"cpu_offload_flag": True}
+        return {}
+
     def load_customized(
-        self, component_model_path: str, server_args: ServerArgs, component_name: str
+        self,
+        component_model_path: str,
+        server_args: ServerArgs,
+        component_name: str,
+        cpu_offload_flag: bool = False,
     ):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_component_config(component_path=component_model_path)
@@ -133,8 +149,9 @@ class VAELoader(ComponentLoader):
             # NOTE: some post init logics are only available after updated with config
             vae_config.post_init()
 
-        component_starts_on_cpu = server_args.should_start_component_on_cpu(
+        component_starts_on_cpu = (
-            component_name
+            server_args.should_start_component_on_cpu(component_name)
+            or cpu_offload_flag
         )
         target_device = self.target_device(component_starts_on_cpu)
 
@@ -190,7 +207,11 @@ class VAELoader(ComponentLoader):
             loaded.update(safetensors_load_file(sf_path))
         _backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
         strict_load = native_only
-        vae.load_state_dict(loaded, strict=strict_load)
+        vae.load_state_dict(
+            loaded,
+            strict=strict_load,
+            assign=bool(cpu_offload_flag and current_platform.is_mps()),
+        )
 
         if not strict_load:
             state_keys = set(vae.state_dict().keys())

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -295,6 +295,15 @@ def maybe_load_fsdp_model(
         logger.info("Disabling FSDP for MPS platform as it's not compatible")
 
     weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device)
+    mps_zero_copy_weight_loading = bool(
+        current_platform.is_mps()
+        and weight_load_plan.mps_layerwise_cpu_staging
+        and weight_load_plan.checkpoint_load_device.type == "cpu"
+    )
+    if mps_zero_copy_weight_loading:
+        # layerwise offload replaces block parameters with mps placeholders after
+        # load, so compatible checkpoint tensors stay file-backed on CPU
+        model._mps_zero_copy_weight_loading = True
     defer_cpu_placement = bool(
         component_starts_on_cpu
         and weight_load_plan.defer_cpu_placement
@@ -408,6 +417,7 @@ def maybe_load_fsdp_model(
         strict=strict,
         cpu_offload=load_on_cpu,
         param_names_mapping=param_names_mapping_fn,
+        mps_zero_copy_weight_loading=mps_zero_copy_weight_loading,
         preconverted_state_dict=preconverted_state_dict,
     )
     if bnb_quant_states:
@@ -531,6 +541,7 @@ def load_model_from_full_model_state_dict(
     strict: bool = False,
     cpu_offload: bool = False,
     param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
+    mps_zero_copy_weight_loading: bool = False,
     preconverted_state_dict: (
         tuple[
             dict[
@@ -555,6 +566,7 @@ def load_model_from_full_model_state_dict(
         strict (bool): flag to check if to load the model in strict mode
         cpu_offload (bool): flag to check if FSDP offload is enabled
         param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
+        mps_zero_copy_weight_loading (bool): retain compatible CPU checkpoint tensors for MPS layerwise offload
     Returns:
         ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
             * **missing_keys** is a list of str containing the missing keys
@@ -696,7 +708,17 @@ def load_model_from_full_model_state_dict(
                 if actual_param is not None
                 else None
             )
-            if weight_loader is not None:
+            use_checkpoint_tensor_directly = bool(
+                mps_zero_copy_weight_loading
+                and actual_param is not None
+                and not getattr(actual_param, "mps_zero_copy_unsafe", False)
+                and tuple(meta_sharded_param.shape) == tuple(full_tensor.shape)
+                and full_tensor.device.type == "cpu"
+                and full_tensor.dtype == target_dtype
+            )
+            if use_checkpoint_tensor_directly:
+                sharded_tensor = full_tensor
+            elif weight_loader is not None:
                 assert actual_param is not None
                 if _can_assign_cpu_tensor_without_copy(
                     actual_param,

python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py

@@ -11,6 +11,9 @@ class WeightLoadPlan:
     checkpoint_load_device: torch.device
     # Device required while running process_weights_after_loading; None means unchanged.
     weight_postprocess_device: torch.device | None = None
+    # mps layerwise loading retains compatible safetensors tensors as CPU-backed
+    # parameters instead of materializing a second unified-memory copy
+    mps_layerwise_cpu_staging: bool = False
     # Delay final CPU placement until after device-side weight postprocessing.
     defer_cpu_placement: bool = False
     # keep the complete mapped checkpoint state dict on the load device
@@ -24,6 +27,7 @@ class WeightLoadPlan:
         needs_device_weight_postprocess: bool,
         component_starts_on_cpu: bool,
         load_full_state_dict_on_device: bool = False,
+        mps_layerwise_cpu_staging: bool = False,
     ) -> "WeightLoadPlan":
         # if on-device weight postprocessing is required, load directly to device to speedup loading
         weight_postprocess_device = (
@@ -36,4 +40,5 @@ class WeightLoadPlan:
                 needs_device_weight_postprocess and component_starts_on_cpu
             ),
             load_full_state_dict_on_device=load_full_state_dict_on_device,
+            mps_layerwise_cpu_staging=mps_layerwise_cpu_staging,
         )

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -65,9 +65,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.post_training.gpu_worker_post_training_mixin import (
     GPUWorkerPostTrainingMixin,
 )
-from sglang.multimodal_gen.runtime.realtime.session import (
+from sglang.multimodal_gen.runtime.realtime.session import RealtimeSessionCache
-    RealtimeSessionCache,
-)
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
 from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
@@ -228,7 +226,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
 
     def init_device_and_model(self) -> None:
         """Initialize the device and load the model."""
-        current_platform.set_device(current_platform.get_device(self.local_rank))
+        if not current_platform.is_mps():
+            current_platform.set_device(current_platform.get_device(self.local_rank))
         # num_gpus is the total world size across every node; the co-located,
         # CPU-contending worker count on THIS host is num_gpus // nnodes.
         local_num_gpus = self.server_args.num_gpus // self.server_args.nnodes
@@ -316,9 +315,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
         if output_batch.metrics:
             output_batch.metrics.record_memory_snapshot("mem_analysis", final_snapshot)
 
-        # for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
+        peak_reserved_bytes = final_snapshot.peak_reserved_mb * (1024**2)
-        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
+        peak_allocated_bytes = final_snapshot.peak_allocated_mb * (1024**2)
-        peak_allocated_bytes = torch.get_device_module().max_memory_allocated()
 
         output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
         peak_reserved_gb = peak_reserved_bytes / (1024**3)
@@ -466,7 +464,11 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
         output_batch = None
         forward_failed = False
         try:
-            if self.is_output_rank and not current_platform.is_cpu():
+            if (
+                self.is_output_rank
+                and not current_platform.is_cpu()
+                and not current_platform.is_mps()
+            ):
                 torch.get_device_module().reset_peak_memory_stats()
 
             start_time = (
@@ -690,8 +692,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
     def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
         if not self.is_output_rank or current_platform.is_cpu():
             return
-        peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
+        output_batch.peak_memory_mb = capture_memory_snapshot().peak_reserved_mb
-        output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
 
     def _forward_group(self, batch: list[Req]) -> OutputBatch:
         assert self.pipeline is not None

python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py

@@ -202,6 +202,14 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
         state: ResidencyState,
     ) -> None:
         if isinstance(module, LayerwiseOffloadableModuleMixin):
+            # MPS layerwise components retain checkpoint-backed CPU weights and
+            # synchronously materialize one layer at a time. Moving the whole
+            # module here would defeat that bounded-residency contract.
+            if current_platform.is_mps():
+                if module.mps_stream_non_layer_weights:
+                    return
+                _module_to_local_device(module, dtype=use.target_dtype)
+                return
             module.prepare_for_next_req()
 
     def finish_use(
@@ -214,6 +222,10 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
             return
         for manager in module.layerwise_offload_managers:
             manager.release_all()
+        if current_platform.is_mps():
+            torch.mps.synchronize()
+            module.restore_mps_cpu_non_layer_weights()
+            torch.mps.empty_cache()
 
     def finish_request(
         self,

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -1,6 +1,7 @@
 import bisect
 import re
 from collections.abc import Mapping, Sequence
+from contextlib import nullcontext
 from typing import Any, Dict, List, Set, Tuple
 
 import torch
@@ -84,7 +85,8 @@ class LayerwiseOffloadManager:
     """A lightweight layerwise CPU offload manager.
 
     This utility offloads per-layer parameters/buffers from GPU to CPU, and
-    supports async H2D prefetch using a dedicated CUDA stream.
+    supports async H2D prefetch using a dedicated CUDA stream. MPS uses
+    synchronous per-layer transfers from the checkpoint-backed CPU tensors.
 
     Typical usage:
     - Construct the manager with the target model and the list-like module
@@ -104,13 +106,23 @@ class LayerwiseOffloadManager:
         pin_cpu_memory: bool = True,
         prefetch_size: int = 1,
         resident_layers: int = 0,
+        initialize: bool = True,
         residency_policy: str = RESIDENCY_POLICY_LEADING,
     ) -> None:
         self.model = model
         self.layers_attr_str = layers_attr_str
         self.num_layers = num_layers
-        self.pin_cpu_memory = pin_cpu_memory
+        self._synchronous_mps = current_platform.is_mps()
-        self.prefetch_size = min(max(1, prefetch_size), self.num_layers)
+        # mps shares physical memory with the CPU and has no pinned host memory
+        # or CUDA-style copy streams
+        self.pin_cpu_memory = bool(pin_cpu_memory and not self._synchronous_mps)
+        # an explicit MPS zero avoids staging the next layer alongside the
+        # active one; MPS has no transfer overlap to recover from that cost
+        self.prefetch_size = (
+            0
+            if current_platform.is_mps() and prefetch_size == 0
+            else min(max(1, prefetch_size), self.num_layers)
+        )
         # Layers held on GPU across denoise steps, instead of being re-streamed
         # every step. `residency_policy` picks *which* layers those are; see
         # compute_streamed_layers for why the choice is not cosmetic.
@@ -134,10 +146,17 @@ class LayerwiseOffloadManager:
         self.enabled = bool(enabled and torch.get_device_module().is_available())
         if not self.enabled:
             return
-        self.device = torch.device(
+        self.device = (
-            current_platform.device_type, torch.get_device_module().current_device()
+            current_platform.get_local_torch_device()
+            if current_platform.is_mps()
+            else torch.device(
+                current_platform.device_type,
+                torch.get_device_module().current_device(),
+            )
+        )
+        self.copy_stream = (
+            None if self._synchronous_mps else torch.get_device_module().Stream()
         )
-        self.copy_stream = torch.get_device_module().Stream()
 
         # ``named_parameters()`` is relative to ``model``, just like the path in
         # ``layers_attr_str``. Anchor the match so a manager for top-level
@@ -153,6 +172,9 @@ class LayerwiseOffloadManager:
         # layer_idx -> {name: pinned_cpu_tensor_with_original_stride}
         # stores tensors whose original non-contiguous stride/layout must be preserved
         self._strided_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
+        # mps keeps the original CPU tensor for each layer instead of building a
+        # second flattened host copy
+        self._mps_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
         # layer_idx -> {name: {dtype, offset, numel, shape}}
         # stores the offset and numel of each weight from a same layer, of same dtype
         self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
@@ -168,6 +190,10 @@ class LayerwiseOffloadManager:
         # Store forward hooks for removal
         self._forward_hooks: List[Any] = []
 
+        if initialize:
+            self._initialize()
+
+    def initialize(self) -> None:
         self._initialize()
 
     def _match_layer_idx(self, name: str) -> int | None:
@@ -229,6 +255,10 @@ class LayerwiseOffloadManager:
         self._named_parameters = dict(self.model.named_parameters())
         self._named_buffers = dict(self.model.named_buffers())
 
+        if self._synchronous_mps:
+            self._initialize_mps_cpu_weights()
+            return
+
         # 1. collect and group layer parameters by dtype. Keep buffers resident:
         # shared buffers such as RoPE caches may be referenced by many layers.
         layer_groups: Dict[int, Dict[torch.dtype, List[Tuple[str, torch.Tensor]]]] = {}
@@ -361,6 +391,33 @@ class LayerwiseOffloadManager:
             return list(range(count))
         return self._next_streamed(after=-1, count=count)
 
+    @torch.compiler.disable
+    def _initialize_mps_cpu_weights(self) -> None:
+        for name, tensor in self._named_parameters.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            local_tensor = self._to_local_tensor(tensor).detach()
+            cpu_tensor = (
+                local_tensor
+                if local_tensor.device.type == "cpu"
+                else local_tensor.to("cpu")
+            )
+            self._mps_cpu_weights.setdefault(layer_idx, {})[name] = cpu_tensor
+            self._weight_metadata.setdefault(layer_idx, {})[name] = {
+                "dtype": cpu_tensor.dtype,
+            }
+            tensor.data = self._get_shared_empty_tensor_for_target(
+                tensor, cpu_tensor.dtype
+            )
+
+        torch.mps.empty_cache()
+        self.register_forward_hooks()
+        self._configured = True
+        logger.info(
+            f"Initialized synchronous MPS layerwise offload with {self.num_layers} layers"
+        )
+
     def prepare_for_next_req(self, non_blocking=True):
         """
         Prepare for the next round of denoising loop with prefetching the necessary layers
@@ -435,22 +492,41 @@ class LayerwiseOffloadManager:
         """
         idempotent
         """
-        if not self.enabled or self.device is None or self.copy_stream is None:
+        if not self.enabled or self.device is None:
             return
         if layer_idx < 0 or layer_idx >= self.num_layers:
             return
         if layer_idx in self._gpu_layers:
             return
+        if self._synchronous_mps:
+            cpu_weights = self._mps_cpu_weights.get(layer_idx)
+            if not cpu_weights:
+                return
+            with torch.inference_mode(False), torch.no_grad():
+                for name, cpu_tensor in cpu_weights.items():
+                    target = self.get_target_with_name(name)
+                    target.data = self._wrap_for_target(
+                        target,
+                        cpu_tensor.to(device=self.device, non_blocking=False),
+                    )
+            self._gpu_layers.add(layer_idx)
+            return
         if layer_idx not in self._consolidated_cpu_weights:
             return
-        self.copy_stream.wait_stream(torch.get_device_module().current_stream())
+        if self.copy_stream is not None:
+            self.copy_stream.wait_stream(torch.get_device_module().current_stream())
+            stream_context = torch.get_device_module().stream(self.copy_stream)
+        else:
+            # the device has no CUDA-like stream or pinned-memory support
+            non_blocking = False
+            stream_context = nullcontext()
 
         # create gpu buffer and load from CPU buffer
         gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
         with (
             torch.inference_mode(False),
             torch.no_grad(),
-            torch.get_device_module().stream(self.copy_stream),
+            stream_context,
         ):
             for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items():
                 gpu_buffer = torch.empty(
@@ -488,10 +564,11 @@ class LayerwiseOffloadManager:
                 ].view(meta["shape"])
                 target.data = self._wrap_for_target(target, local_tensor)
 
-        # record the prefetch event of this layer after all copies are enqueued
+        if self.copy_stream is not None:
-        event = torch.get_device_module().Event()
+            # record after all copies so the consumer waits for every weight copy
-        event.record(self.copy_stream)
+            event = torch.get_device_module().Event()
-        self._prefetch_events[layer_idx] = event
+            event.record(self.copy_stream)
+            self._prefetch_events[layer_idx] = event
 
         self._gpu_layers.add(layer_idx)
 
@@ -524,6 +601,11 @@ class LayerwiseOffloadManager:
                 )
 
         self._gpu_layers.discard(layer_idx)
+        if self._synchronous_mps:
+            # mps dispatch is asynchronous, so a tensor rebinding alone leaves
+            # prior layer allocations live until the command buffer drains
+            torch.mps.synchronize()
+            torch.mps.empty_cache()
 
     @torch.compiler.disable
     def release_all(self) -> None:
@@ -554,6 +636,10 @@ class LayerwiseOffloadManager:
         """Sync a layer's weights from GPU back to CPU."""
         if not self.enabled or layer_idx not in self._gpu_layers:
             return
+        if self._synchronous_mps:
+            # inference does not mutate parameters; retain the original mapped
+            # tensor instead of materializing a CPU copy after every layer
+            return
         if layer_idx not in self._consolidated_cpu_weights:
             return
 
@@ -616,6 +702,33 @@ class LayerwiseOffloadManager:
             return None
 
         updated_names: Set[str] = set()
+        if self._synchronous_mps:
+            for name, loaded_weight in weight_dict.items():
+                layer_idx = self._match_layer_idx(name)
+                if layer_idx is None:
+                    continue
+                cpu_tensor = self._mps_cpu_weights.get(layer_idx, {}).get(name)
+                if cpu_tensor is None:
+                    continue
+                if tuple(cpu_tensor.shape) != tuple(loaded_weight.shape):
+                    raise ValueError(
+                        f"Shape mismatch for {name}: "
+                        f"expected={tuple(cpu_tensor.shape)}, "
+                        f"loaded={tuple(loaded_weight.shape)}"
+                    )
+                replacement = loaded_weight.to(
+                    device="cpu", dtype=cpu_tensor.dtype
+                ).detach()
+                self._mps_cpu_weights[layer_idx][name] = replacement
+                if layer_idx in self._gpu_layers:
+                    target = self.get_target_with_name(name)
+                    target.data = self._wrap_for_target(
+                        target,
+                        replacement.to(device=target.device, dtype=target.dtype),
+                    )
+                updated_names.add(name)
+            return updated_names
+
         for name, loaded_weight in weight_dict.items():
             layer_idx = self._match_layer_idx(name)
             if layer_idx is None:
@@ -665,6 +778,11 @@ class LayerwiseOffloadManager:
         when offload is enabled, this method returns the real weights and
         can be used for checksum computation.
         """
+        if self._synchronous_mps:
+            for layer_idx in sorted(self._mps_cpu_weights):
+                yield from self._mps_cpu_weights[layer_idx].items()
+            return
+
         for layer_idx in sorted(self._weight_metadata):
             for name, meta in self._weight_metadata[layer_idx].items():
                 if meta.get("preserve_strides", False):
@@ -696,7 +814,7 @@ class LayerwiseOffloadManager:
                 if i not in self._gpu_layers:
                     # LTX audio VAE traverses decoder.up in reverse order
                     self.prefetch_layer(i, non_blocking=False)
-                if i in self._prefetch_events:
+                if i in self._prefetch_events and self.copy_stream is not None:
                     torch.get_device_module().current_stream().wait_event(
                         self._prefetch_events[i]
                     )
@@ -716,7 +834,7 @@ class LayerwiseOffloadManager:
                     ):
                         self.prefetch_layer(layer_to_prefetch, non_blocking=True)
                 # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed
-                elif i % self.prefetch_size == 0:
+                elif self.prefetch_size and i % self.prefetch_size == 0:
                     for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size):
                         layer_to_prefetch = j % self.num_layers
                         self.prefetch_layer(layer_to_prefetch, non_blocking=True)
@@ -750,11 +868,106 @@ class LayerwiseOffloadableModuleMixin:
 
     # whether the current module is selected by the `dit` group
     layerwise_offload_dit_group_enabled: bool = True
+    # H3 has a large packed-sequence working set on MPS, so its non-block
+    # weights are materialized only for the subphase that consumes them
+    mps_stream_non_layer_weights: bool = False
 
     # The list of names of this module's layer/block ModuleList or Sequential attributes.
     layer_names: List[str] = []
     layerwise_offload_managers: list[LayerwiseOffloadManager] = []
 
+    def _capture_mps_cpu_non_layer_weights(self) -> None:
+        managed_names = {
+            name
+            for manager in self.layerwise_offload_managers
+            for weights in manager._mps_cpu_weights.values()
+            for name in weights
+        }
+        self._mps_cpu_non_layer_parameters = {
+            name: parameter.detach()
+            for name, parameter in self.named_parameters()
+            if name not in managed_names
+        }
+        self._mps_cpu_buffers = {
+            name: buffer.detach() for name, buffer in self.named_buffers()
+        }
+        if not self.mps_stream_non_layer_weights:
+            return
+
+        parameters = dict(self.named_parameters())
+        for name, tensor in self._mps_cpu_non_layer_parameters.items():
+            parameters[name].data = torch.empty(
+                (1,),
+                dtype=tensor.dtype,
+                device=current_platform.get_local_torch_device(),
+            )
+        buffers = dict(self.named_buffers())
+        for name, tensor in self._mps_cpu_buffers.items():
+            buffers[name].data = torch.empty(
+                (1,),
+                dtype=tensor.dtype,
+                device=current_platform.get_local_torch_device(),
+            )
+
+    @staticmethod
+    def _matches_mps_weight_prefix(name: str, prefixes: tuple[str, ...]) -> bool:
+        return any(
+            name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes
+        )
+
+    def materialize_mps_non_layer_weights(self, *prefixes: str) -> None:
+        if not current_platform.is_mps() or not self.mps_stream_non_layer_weights:
+            return
+        selected_prefixes = tuple(prefixes)
+        with torch.inference_mode(False), torch.no_grad():
+            parameters = dict(self.named_parameters())
+            for name, tensor in self._mps_cpu_non_layer_parameters.items():
+                if self._matches_mps_weight_prefix(name, selected_prefixes):
+                    parameters[name].data = tensor.to(
+                        current_platform.get_local_torch_device()
+                    )
+            buffers = dict(self.named_buffers())
+            for name, tensor in self._mps_cpu_buffers.items():
+                if self._matches_mps_weight_prefix(name, selected_prefixes):
+                    buffers[name].data = tensor.to(
+                        current_platform.get_local_torch_device()
+                    )
+
+    def release_mps_non_layer_weights(self, *prefixes: str) -> None:
+        if not current_platform.is_mps() or not self.mps_stream_non_layer_weights:
+            return
+        selected_prefixes = tuple(prefixes)
+        with torch.inference_mode(False), torch.no_grad():
+            parameters = dict(self.named_parameters())
+            for name, tensor in self._mps_cpu_non_layer_parameters.items():
+                if self._matches_mps_weight_prefix(name, selected_prefixes):
+                    parameters[name].data = torch.empty(
+                        (1,),
+                        dtype=tensor.dtype,
+                        device=current_platform.get_local_torch_device(),
+                    )
+            buffers = dict(self.named_buffers())
+            for name, tensor in self._mps_cpu_buffers.items():
+                if self._matches_mps_weight_prefix(name, selected_prefixes):
+                    buffers[name].data = torch.empty(
+                        (1,),
+                        dtype=tensor.dtype,
+                        device=current_platform.get_local_torch_device(),
+                    )
+        torch.mps.synchronize()
+        torch.mps.empty_cache()
+
+    def restore_mps_cpu_non_layer_weights(self) -> None:
+        if not current_platform.is_mps():
+            return
+        with torch.inference_mode(False), torch.no_grad():
+            parameters = dict(self.named_parameters())
+            for name, tensor in self._mps_cpu_non_layer_parameters.items():
+                parameters[name].data = tensor
+            buffers = dict(self.named_buffers())
+            for name, tensor in self._mps_cpu_buffers.items():
+                buffers[name].data = tensor
+
     def configure_layerwise_offload(self, server_args: ServerArgs):
         self.layerwise_offload_managers = []
         named_modules = dict(self.named_modules())
@@ -774,7 +987,9 @@ class LayerwiseOffloadableModuleMixin:
             prefetch_value = (
                 server_args.dit_offload_prefetch_size if dit_tuning_enabled else 0.0
             )
-            if prefetch_value < 1.0:
+            if current_platform.is_mps() and prefetch_value == 0.0:
+                prefetch_size = 0
+            elif prefetch_value < 1.0:
                 prefetch_size = 1 + int(round(prefetch_value * (num_layers - 1)))
             else:
                 prefetch_size = int(prefetch_value)
@@ -797,6 +1012,7 @@ class LayerwiseOffloadableModuleMixin:
                 pin_cpu_memory=server_args.pin_cpu_memory,
                 prefetch_size=prefetch_size,
                 resident_layers=resident_layers,
+                initialize=not current_platform.is_mps(),
                 residency_policy=(
                     server_args.dit_layerwise_residency_policy
                     if dit_tuning_enabled
@@ -806,6 +1022,11 @@ class LayerwiseOffloadableModuleMixin:
             self.layerwise_offload_managers.append(manager)
             configured_layer_names.append(layer_name)
 
+        if current_platform.is_mps():
+            for manager in self.layerwise_offload_managers:
+                manager.initialize()
+            self._capture_mps_cpu_non_layer_weights()
+
         if configured_layer_names:
             logger.debug(
                 "Enabled layerwise offload for %s on modules: %s",

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -77,6 +77,19 @@ logger = init_logger(__name__)
 _ARCH_DEFAULTS = MiniMaxH3DiTArchConfig()
 _BF16_DTYPE = torch.bfloat16
 _FP32_DTYPE = torch.float32
+_MPS_MLP_TOKEN_CHUNK_SIZE = 128
+# keep MPS activation chunks below the allocator high-watermark; CUDA keeps
+# its fused full-sequence projection
+_MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE = 128
+_MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE = 128
+
+_MPS_EMBED_WEIGHT_PREFIXES = (
+    "condition_proj",
+    "video_patch_proj",
+    "audio_patch_proj",
+    "time_embedder",
+    "token_refiner.final_norm",
+)
 
 _MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER = (
     "video_patch_proj.weight",
@@ -390,6 +403,18 @@ def _apply_rope_qk(
     return q, k
 
 
+def _apply_rope(
+    x: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+) -> torch.Tensor:
+    """Apply the eager (non-CUDA) H3 RoPE path to one Q or K tensor."""
+    half = cos_sin_cache.shape[-1] // 2
+    cos_half, sin_half = cos_sin_cache.split(half, dim=-1)
+    cos = torch.cat((cos_half, cos_half), dim=-1).unsqueeze(1)
+    sin = torch.cat((sin_half, sin_half), dim=-1).unsqueeze(1)
+    return _apply_rope_cos_sin(x, cos, sin)
+
+
 class MiniMaxH3TimeEmbedder(nn.Module):
     def __init__(
         self,
@@ -597,6 +622,9 @@ class MiniMaxH3Attention(nn.Module):
 
     def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None:
         weight = self.qkv_proj.weight
+        # h3 checkpoints interleave each attention head's Q, K, and V rows
+        # this parameter needs reordering before the native QKV projection
+        weight.mps_zero_copy_unsafe = True
         base_loader = weight.weight_loader
 
         def _reorder_checkpoint_weight(loaded_weight: torch.Tensor) -> torch.Tensor:
@@ -630,6 +658,107 @@ class MiniMaxH3Attention(nn.Module):
         # rank-local FSDP must reorder grouped QKV before selecting each shard
         weight.rank_local_weight_transform = _reorder_checkpoint_weight
 
+    def _forward_mps_streamed_attention(
+        self,
+        x: torch.Tensor,
+        *,
+        rope_cache: tuple[torch.Tensor, torch.Tensor] | None,
+        cu_seqlens: torch.Tensor,
+        cu_seqlens_host: tuple[int, ...] | None,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        """Run MPS attention without materializing the full QKV activation.
+
+        H3's fused QKV output alone is roughly 1.45 GiB at 768px.  MPS shares
+        unified memory with the host, so retaining it alongside the packed
+        residual and SDPA workspace can evict the OS.  Build normalized K/V
+        once, then project Q a small chunk at a time and immediately consume it
+        through attention and the output projection.  The formula, weights,
+        and complete K/V context are unchanged; this is intentionally limited
+        to single-device MPS where Ulysses collectives are not active.
+        """
+        total = x.shape[0]
+        key = torch.empty(
+            (total, self.num_heads, self.head_dim), dtype=x.dtype, device=x.device
+        )
+        value = torch.empty_like(key)
+        cos_sin_cache = None if rope_cache is None else rope_cache[0]
+
+        # Do not retain Q while producing the K/V cache.  Dropping the chunk
+        # before the next transfer keeps only two full-width attention tensors.
+        for start in range(0, total, _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE):
+            stop = min(start + _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE, total)
+            qkv, _ = self.qkv_proj(x[start:stop])
+            q_chunk, k_chunk, v_chunk = qkv.split(self.local_inner_dim, dim=-1)
+            del q_chunk
+            k_chunk = self.k_norm(k_chunk.view(-1, self.num_heads, self.head_dim))
+            if cos_sin_cache is not None:
+                k_chunk = _apply_rope(k_chunk, cos_sin_cache[start:stop])
+            key[start:stop].copy_(k_chunk)
+            value[start:stop].copy_(v_chunk.view(-1, self.num_heads, self.head_dim))
+            del qkv, k_chunk, v_chunk
+            torch.mps.synchronize()
+            torch.mps.empty_cache()
+
+        if self._attention_impl is None:
+            self._set_attention_backend(
+                get_attn_backend(
+                    self.head_dim,
+                    x.dtype,
+                    attention_requirements=AttentionRequirements(packed_varlen=True),
+                )
+            )
+        bounds = (
+            cu_seqlens_host
+            if cu_seqlens_host is not None
+            else tuple(int(item) for item in cu_seqlens.tolist())
+        )
+        out = torch.empty_like(x)
+        for sequence_start, sequence_stop in zip(bounds[:-1], bounds[1:]):
+            if sequence_start == sequence_stop:
+                continue
+            keys = key[sequence_start:sequence_stop].unsqueeze(0)
+            values = value[sequence_start:sequence_stop].unsqueeze(0)
+            for start in range(
+                sequence_start,
+                sequence_stop,
+                _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE,
+            ):
+                stop = min(start + _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE, sequence_stop)
+                qkv, _ = self.qkv_proj(x[start:stop])
+                q_chunk, k_chunk, v_chunk = qkv.split(self.local_inner_dim, dim=-1)
+                del k_chunk, v_chunk
+                for query_start in range(
+                    start, stop, _MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE
+                ):
+                    query_stop = min(
+                        query_start + _MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE, stop
+                    )
+                    q = self.q_norm(
+                        q_chunk[query_start - start : query_stop - start].view(
+                            -1, self.num_heads, self.head_dim
+                        )
+                    )
+                    if cos_sin_cache is not None:
+                        q = _apply_rope(q, cos_sin_cache[query_start:query_stop])
+                    attention_out = self._attention_impl.forward(
+                        q.unsqueeze(0), keys, values, None
+                    )[0]
+                    projected, _ = self.out_proj(
+                        attention_out.reshape(
+                            query_stop - query_start, self.local_inner_dim
+                        )
+                    )
+                    out[query_start:query_stop].copy_(projected)
+                    del q, attention_out, projected
+                del qkv, q_chunk
+                torch.mps.synchronize()
+                torch.mps.empty_cache()
+
+        del key, value
+        torch.mps.empty_cache()
+        return out
+
     def forward(
         self,
         x: torch.Tensor,
@@ -652,6 +781,15 @@ class MiniMaxH3Attention(nn.Module):
         so cu_seqlens retains global packed-document semantics. The inverse
         all-to-all restores the row shard before the output projection.
         """
+        if x.device.type == "mps" and not ulysses_active:
+            return self._forward_mps_streamed_attention(
+                x,
+                rope_cache=rope_cache,
+                cu_seqlens=cu_seqlens,
+                cu_seqlens_host=cu_seqlens_host,
+                max_seqlen=max_seqlen,
+            )
+
         total = x.shape[0]
         qkv, _ = self.qkv_proj(x)
         q, k, v = qkv.split(self.local_inner_dim, dim=-1)
@@ -747,6 +885,18 @@ class MiniMaxH3MLP(nn.Module):
         self.reuse_fc1_activation = quant_config is None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.device.type == "mps":
+            out = torch.empty_like(x)
+            for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
+                stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
+                hidden, _ = self.fc1(x[start:stop])
+                hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
+                chunk, _ = self.fc2(hidden)
+                out[start:stop].copy_(chunk)
+                del hidden, chunk
+                torch.mps.synchronize()
+                torch.mps.empty_cache()
+            return out
         hidden, _ = self.fc1(x)
         hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
         out, _ = self.fc2(hidden)
@@ -1355,6 +1505,38 @@ class MiniMaxH3FinalLayer(nn.Module):
                 raise ValueError("MiniMax H3 AdaLN cache parameters are required")
             adaln_params = self.adaln_proj(adaln_input)
         shift, scale = adaln_params
+        if x.device.type == "mps":
+            video = audio = None
+            for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
+                stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
+                h = self.norm(x[start:stop])
+                h = _modulate_scale_shift(
+                    h,
+                    shift,
+                    scale,
+                    inverse_indices[start:stop],
+                    dtype=_BF16_DTYPE,
+                ).to(_FP32_DTYPE)
+                video_chunk, _ = self.video_out(h)
+                audio_chunk, _ = self.audio_out(h)
+                if video is None:
+                    video = torch.empty(
+                        (x.shape[0], video_chunk.shape[-1]),
+                        dtype=video_chunk.dtype,
+                        device=x.device,
+                    )
+                    audio = torch.empty(
+                        (x.shape[0], audio_chunk.shape[-1]),
+                        dtype=audio_chunk.dtype,
+                        device=x.device,
+                    )
+                video[start:stop].copy_(video_chunk)
+                audio[start:stop].copy_(audio_chunk)
+                del h, video_chunk, audio_chunk
+                torch.mps.synchronize()
+                torch.mps.empty_cache()
+            assert video is not None and audio is not None
+            return video, audio
         h = self.norm(x)
         h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
         # Preserve full precision through both final output projections.
@@ -1371,6 +1553,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
     # parameters mix fp32 (patch projections, timestep embedder, and output
     # heads) with bf16 blocks; FSDP must gather in each parameter's own dtype
     _fsdp_mixed_dtype_params = True
+    mps_stream_non_layer_weights = True
     _compile_conditions = [is_block]
     param_names_mapping = _ARCH_DEFAULTS.param_names_mapping
     reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping
@@ -1551,7 +1734,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                 for index in range(arch.num_layers)
             ]
         )
-        self.layer_names = ["blocks"]
+        self.layer_names = ["token_refiner.blocks", "blocks"]
         self.final_layer = MiniMaxH3FinalLayer(
             arch,
             quant_config,
@@ -1655,6 +1838,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         device: torch.device,
     ) -> torch.Tensor:
         """Project and refine request-static text conditioning once."""
+        self.materialize_mps_non_layer_weights(
+            "condition_proj", "token_refiner.final_norm"
+        )
         text_len = int(refiner_cu_seqlens[1].item())
         if text_len <= 0 or text_len > int(prompt_embeds.shape[0]):
             raise ValueError(
@@ -1670,12 +1856,14 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
             )
         )
         text_embed, _ = self.condition_proj(text_rows)
-        return self.token_refiner(
+        refined = self.token_refiner(
             text_embed,
             cu_seqlens=true_refiner_cu,
             cu_seqlens_host=(0, text_len, text_len),
             max_seqlen=text_len,
         )
+        self.release_mps_non_layer_weights("condition_proj", "token_refiner.final_norm")
+        return refined
 
     def build_rope_cache(
         self,
@@ -1690,6 +1878,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         chunk) -- see forward()'s row_start derivation for the identity
         this must stay in sync with.
         """
+        self.materialize_mps_non_layer_weights("rope")
         if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1:
             raise ValueError(
                 "img_position_ids must be [1, S, 3], got "
@@ -1711,7 +1900,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         rope_freqs = self.rope(
             img_position_ids[:, row_start : row_start + local_seq_len]
         ).to(device)
-        return (
+        result = (
             _rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
             torch.arange(
                 local_seq_len,
@@ -1719,6 +1908,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                 dtype=torch.long,
             ),
         )
+        self.release_mps_non_layer_weights("rope")
+        return result
 
     @eager_on_graph(True)
     def _embed(
@@ -1994,6 +2185,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         # request-static cache once; direct model callers use this fallback.
         rope_cache = kwargs.get("rope_cache")
         if rope_cache is None:
+            self.materialize_mps_non_layer_weights("rope")
             rope_freqs = self.rope(img_position_ids[:, row_start:row_stop]).to(device)
             rope_cache = (
                 _rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
@@ -2003,6 +2195,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                     dtype=torch.long,
                 ),
             )
+            self.release_mps_non_layer_weights("rope")
+        self.materialize_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES)
         img_pos = img_pos.to(device)
         audio_pos = audio_pos.to(device)
         text_pos = text_pos.to(device)
@@ -2023,6 +2217,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
             refined_prompt_embeds_length=kwargs.get("refined_prompt_embeds_length"),
             local_embedding_layout=kwargs.get("local_embedding_layout"),
         )
+        self.release_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES)
         # request-step AdaLN input shared by all blocks
         adaln_input = nn.functional.silu(t_emb).to(_BF16_DTYPE)
         inverse_indices = inverse_indices.to(device)
@@ -2091,6 +2286,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                     None if block_adaln_params is None else block_adaln_params[index]
                 ),
             )
+        self.materialize_mps_non_layer_weights("final_layer")
         video_logits, audio_logits = self.final_layer(
             hidden,
             adaln_input=adaln_input,
@@ -2104,6 +2300,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                 )
             ),
         )
+        self.release_mps_non_layer_weights("final_layer")
         if sp_ws > 1:
             from sglang.multimodal_gen.runtime.distributed.parallel_state import (
                 get_sp_group,

python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py

@@ -41,9 +41,10 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
     eight otherwise-idle ranks during encoding.
     """
 
-    layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
-
     supports_dp_encode = True
+    # The inherited text-layer list covers Qwen's language stack; reference
+    # modes also execute the embedded visual tower.
+    layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
     supported_checkpoint_quantization_methods = frozenset({"fp8"})
 
     @staticmethod
@@ -191,7 +192,18 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
                 )
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
             try:
-                weight_loader(param, loaded_weight.to(param.dtype))
+                can_keep_checkpoint_tensor = bool(
+                    getattr(self, "_mps_zero_copy_weight_loading", False)
+                    and weight_loader is default_weight_loader
+                    and param.device.type == "cpu"
+                    and loaded_weight.device.type == "cpu"
+                    and loaded_weight.dtype == param.dtype
+                    and tuple(loaded_weight.shape) == tuple(param.shape)
+                )
+                if can_keep_checkpoint_tensor:
+                    param.data = loaded_weight
+                else:
+                    weight_loader(param, loaded_weight.to(param.dtype))
             except Exception as exc:
                 raise RuntimeError(
                     "Failed to load MiniMax H3 Qwen3-VL weight "

python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Transformer building blocks for the MiniMax H3 visual VAE ViT decoder.
 import math
+from contextlib import nullcontext
 from typing import Optional
 
 import torch
@@ -171,7 +172,10 @@ class RotaryEmbeddingND(nn.Module):
         if D != self.n_dim:
             raise ValueError(f"Expected {self.n_dim} dimensions, got {D}")
 
-        with torch.autocast("cuda", enabled=False):
+        autocast_context = (
+            torch.autocast("cuda", enabled=False) if img_ids.is_cuda else nullcontext()
+        )
+        with autocast_context:
             angles = (
                 self.angle_scale
                 * img_ids[:, :, :, None]

python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle).
+from contextlib import nullcontext
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -22,6 +24,10 @@ def _linear_with_module_dtype(linear, tensor, out_dtype=None):
     return output
 
 
+def _cuda_autocast_disabled(tensor: torch.Tensor):
+    return torch.autocast("cuda", enabled=False) if tensor.is_cuda else nullcontext()
+
+
 def _pack_tensors_3d(tensors, patch_size, patch_size_t):
     batch_size, num_channels_tensors, temporal, height, width = tensors.shape
 
@@ -264,7 +270,7 @@ class ViT3DDecoder(ViTBase):
         hidden_states = _pack_tensors_3d(x, 1, 1)
         latent_size = (latent_T, latent_H, latent_W)
 
-        with torch.autocast("cuda", enabled=False):
+        with _cuda_autocast_disabled(hidden_states):
             hidden_states = _linear_with_module_dtype(
                 self.x_embedder, hidden_states, hidden_states.dtype
             )
@@ -339,7 +345,7 @@ class ViT3DDecoder(ViTBase):
 
         hidden_states = self.apply_mask_postprocess(hidden_states, num_patches)
 
-        with torch.autocast("cuda", enabled=False):
+        with _cuda_autocast_disabled(hidden_states):
             output = _linear_with_module_dtype(
                 self.proj_out, hidden_states, hidden_states.dtype
             )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -1636,7 +1636,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
 
         # deallocate transformer if on mps
         pipeline = self.pipeline() if self.pipeline else None
-        if torch.backends.mps.is_available() and not is_warmup:
+        if (
+            torch.backends.mps.is_available()
+            and not is_warmup
+            and not is_layerwise_offloaded_module(self.transformer)
+        ):
             logger.info(
                 "Memory before deallocating transformer: %s",
                 torch.mps.current_allocated_memory(),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py

@@ -70,6 +70,8 @@ class _AudioVAEDeterminismContext:
     _saved: tuple | None = None
 
     def __enter__(self):
+        if not torch.cuda.is_available():
+            return self
         if _AudioVAEDeterminismContext._depth == 0:
             b = torch.backends
             _AudioVAEDeterminismContext._saved = (
@@ -94,6 +96,8 @@ class _AudioVAEDeterminismContext:
         return self
 
     def __exit__(self, exc_type, exc, tb):
+        if not torch.cuda.is_available():
+            return
         _AudioVAEDeterminismContext._depth -= 1
         if _AudioVAEDeterminismContext._depth == 0:
             b = torch.backends

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py

@@ -147,6 +147,11 @@ class MiniMaxH3PartitionAdmissionStage(PipelineStage):
         if not isinstance(task, str) or not task.strip():
             raise ValueError("MiniMax H3 request task must be a non-empty string")
         self.metadata.canonical_task(task)
+        if batch.num_inference_steps < 2:
+            raise ValueError(
+                "MiniMax H3 requires num_inference_steps >= 2 because its "
+                "video/audio sigma schedules include both interval endpoints"
+            )
         quality = getattr(batch.sampling_params, "quality", "lossless")
         if quality not in QUALITY_LEVELS:
             raise ValueError(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 import functools
 from collections.abc import Mapping
+from contextlib import nullcontext
 
 import torch
 
@@ -296,11 +297,16 @@ class MiniMaxH3DecodingStage(DecodingStage):
                 audio_latent.device.type == "cuda"
                 and autocast_enabled(audio_vae_dtype, server_args.disable_autocast)
             )
-            with torch.autocast(
+            autocast_context = (
-                device_type=audio_latent.device.type,
+                torch.autocast(
-                dtype=audio_vae_dtype,
+                    device_type="cuda",
-                enabled=audio_autocast_enabled,
+                    dtype=audio_vae_dtype,
-            ):
+                    enabled=audio_autocast_enabled,
+                )
+                if audio_latent.is_cuda
+                else nullcontext()
+            )
+            with autocast_context:
                 audio_decode = self._get_vae_decode_fn(
                     audio_vae,
                     server_args,
@@ -352,11 +358,16 @@ class MiniMaxH3DecodingStage(DecodingStage):
             )
             if visual_autocast_enabled:
                 selected_video_vae.prepare_decoder_autocast_weights(video_vae_dtype)
-            with torch.autocast(
+            autocast_context = (
-                device_type=visual_latent.device.type,
+                torch.autocast(
-                dtype=video_vae_dtype,
+                    device_type="cuda",
-                enabled=visual_autocast_enabled,
+                    dtype=video_vae_dtype,
-            ):
+                    enabled=visual_autocast_enabled,
+                )
+                if visual_latent.is_cuda
+                else nullcontext()
+            )
+            with autocast_context:
                 video_decode = self._get_vae_decode_fn(
                     selected_video_vae,
                     server_args,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py

@@ -35,6 +35,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
@@ -605,9 +606,9 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
 
         ctx = _resolve_full_loop_context(batch)
 
-        if not torch.cuda.is_available():
+        if not (current_platform.is_cuda() or current_platform.is_mps()):
-            raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA")
+            raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA or MPS")
-        device = torch.device("cuda")
+        device = current_platform.get_local_torch_device()
         sigmas_video = [float(v) for v in ctx.sigmas["video"]]
         self._maybe_enable_cache_dit_and_torch_compile(
             len(sigmas_video) - 1,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m
 from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
     TextEncodingStage,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -55,6 +56,8 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
             try:
                 self._encode_from_plan(batch, plan)
                 self._publish_native_text_conditioning(batch)
+                if current_platform.is_mps():
+                    self._finish_active_component_use()
             except Exception:
                 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
                     minimax_h3_cleanup_temp_dirs,

python/sglang/multimodal_gen/runtime/platforms/mps.py

@@ -46,15 +46,15 @@ class MpsPlatform(Platform):
 
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
-        raise NotImplementedError
+        return None
 
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
-        raise NotImplementedError
+        return "Apple Silicon MPS"
 
     @classmethod
     def get_device_uuid(cls, device_id: int = 0) -> str:
-        raise NotImplementedError
+        return "mps"
 
     @classmethod
     @lru_cache(maxsize=1)

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -1310,31 +1310,17 @@ class ServerArgs(DisaggServerArgsMixin):
 
     def _adjust_platform_specific(self):
         if current_platform.is_mps():
+            if self.num_gpus != 1:
+                raise ValueError("MPS currently supports only --num-gpus 1")
             if self.component_residency is not None and any(
-                mode != RESIDENT for mode in self.component_residency.values()
+                mode not in (RESIDENT, LAYERWISE_OFFLOAD)
+                for mode in self.component_residency.values()
             ):
                 raise ValueError(
-                    "--component-residency offload modes require CUDA; "
+                    "MPS supports only resident or layerwise-offload component "
-                    "MPS supports only resident components"
+                    "residency"
                 )
             self.use_fsdp_inference = False
-            self.dit_layerwise_offload = False
-            self.layerwise_offload_components = None
-            if (
-                self.dit_cpu_offload
-                or self.text_encoder_cpu_offload
-                or self.image_encoder_cpu_offload
-                or self.vae_cpu_offload
-            ):
-                logger.warning(
-                    "Disabling component CPU offload on MPS because the component "
-                    "residency offload strategy is only validated on CUDA."
-                )
-            self.dit_cpu_offload = False
-            self.text_encoder_cpu_offload = False
-            self.image_encoder_cpu_offload = False
-            self.vae_cpu_offload = False
-            self.cpu_offload_components = None
 
     def is_arg_explicitly_set(self, arg_name: str) -> bool:
         return arg_name in self._explicit_arg_names

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -925,6 +925,7 @@ def maybe_download_model(
         local_path = snapshot_download(
             repo_id=model_name_or_path,
             ignore_patterns=["*.onnx", "*.msgpack"],
+            allow_patterns=allow_patterns,
             local_dir=local_dir,
             local_files_only=True,
             max_workers=8,
@@ -1034,6 +1035,7 @@ def maybe_download_model(
                     local_path = snapshot_download(
                         repo_id=model_name_or_path,
                         ignore_patterns=["*.onnx", "*.msgpack"],
+                        allow_patterns=allow_patterns,
                         local_dir=local_dir,
                         max_workers=8,
                         force_download=True,
@@ -1084,6 +1086,28 @@ def maybe_download_model(
                 wait_time,
             )
             time.sleep(wait_time)
+        except RuntimeError as e:
+            if "client has been closed" not in str(e).lower():
+                raise ValueError(
+                    f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()}: {e}"
+                ) from e
+            if attempt == MAX_RETRIES - 1:
+                raise ValueError(
+                    f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()} "
+                    f"after {MAX_RETRIES} attempts due to network error: {e}"
+                ) from e
+            from huggingface_hub.utils._http import close_session
+
+            close_session()
+            wait_time = 2**attempt
+            logger.warning(
+                "Download failed (attempt %d/%d) because the Hugging Face client was closed. "
+                "Retrying in %d seconds...",
+                attempt + 1,
+                MAX_RETRIES,
+                wait_time,
+            )
+            time.sleep(wait_time)
         except Exception as e:
             raise ValueError(
                 f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()}: {e}"

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -134,6 +134,16 @@ def capture_memory_snapshot() -> MemorySnapshot:
             peak_reserved_mb=0.0,
         )
 
+    if current_platform.is_mps():
+        allocated = torch.mps.current_allocated_memory()
+        reserved = torch.mps.driver_allocated_memory()
+        return MemorySnapshot(
+            allocated_mb=allocated / (1024**2),
+            reserved_mb=reserved / (1024**2),
+            peak_allocated_mb=allocated / (1024**2),
+            peak_reserved_mb=reserved / (1024**2),
+        )
+
     allocated = torch.get_device_module().memory_allocated()
     reserved = torch.get_device_module().memory_reserved()
     peak_allocated = torch.get_device_module().max_memory_allocated()

python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py

@@ -19,6 +19,10 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionRequirements,
 )
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
+    LAYERWISE_OFFLOAD,
+    RESIDENT,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
     MiniMaxH3PartitionAdmissionStage,
     MiniMaxH3ReleaseMetadata,
@@ -369,3 +373,30 @@ def test_validate_server_args_requires_packed_varlen_backend():
     ):
         with pytest.raises(ValueError, match="does not implement packed varlen"):
             MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
+
+
+def test_mps_admission_requires_layerwise_residency_for_every_h3_component():
+    config = SimpleNamespace(
+        vae_config=SimpleNamespace(resolved_parallel_decode_mode=lambda: None),
+        dit_config=SimpleNamespace(arch_config=SimpleNamespace(attention_head_dim=128)),
+        _server_arg_value=MiniMaxH3PipelineConfig._server_arg_value,
+    )
+    modes = {
+        "transformer": LAYERWISE_OFFLOAD,
+        "text_encoder": LAYERWISE_OFFLOAD,
+        "video_vae": LAYERWISE_OFFLOAD,
+        "audio_vae": LAYERWISE_OFFLOAD,
+    }
+    server_args = SimpleNamespace(
+        component_attention_backends={},
+        attention_backend=None,
+        enable_torch_compile=False,
+        residency_mode=modes.get,
+    )
+
+    with patch.object(current_platform, "is_mps", return_value=True):
+        MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
+
+        modes["audio_vae"] = RESIDENT
+        with pytest.raises(ValueError, match="audio_vae"):
+            MiniMaxH3PipelineConfig.validate_server_args(config, server_args)