From 593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 13:35:01 +0800 Subject: [PATCH] [diffusion] optimization: reduce minimax h3 mps memory pressure (#33880) --- .../configs/pipeline_configs/minimax_h3.py | 26 ++ .../runtime/layers/attention/backends/sdpa.py | 34 ++- .../multimodal_gen/runtime/layers/linear.py | 10 +- .../component_loaders/text_encoder_loader.py | 13 +- .../component_loaders/transformer_loader.py | 38 ++- .../loader/component_loaders/vae_loader.py | 29 +- .../runtime/loader/fsdp_load.py | 24 +- .../runtime/loader/weight_load_plan.py | 5 + .../runtime/managers/gpu_worker.py | 21 +- .../component_residency_strategies.py | 12 + .../memory_managers/layerwise_offload.py | 253 ++++++++++++++++-- .../runtime/models/dits/minimax_h3.py | 203 +++++++++++++- .../models/encoders/minimax_h3_qwen3vl.py | 18 +- .../vaes/minimax_h3_video_vae/base_module.py | 6 +- .../vaes/minimax_h3_video_vae/vae_vit.py | 10 +- .../pipelines_core/stages/denoising.py | 6 +- .../minimax_h3/reference_encoding.py | 4 + .../minimax_h3/release_metadata.py | 5 + .../minimax_h3/stages/decoding.py | 31 ++- .../minimax_h3/stages/denoising.py | 7 +- .../minimax_h3/stages/text_encoding.py | 3 + .../multimodal_gen/runtime/platforms/mps.py | 6 +- .../runtime/server_args/server_args.py | 26 +- .../runtime/utils/hf_diffusers_utils.py | 24 ++ .../runtime/utils/perf_logger.py | 10 + .../test/unit/test_minimax_h3_admission.py | 31 +++ 26 files changed, 757 insertions(+), 98 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py index 565f1872f..4c199cff6 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 693dadc62..02bbd5e6b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/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( - query[start:stop].unsqueeze(0), - key[start:stop].unsqueeze(0), - value[start:stop].unsqueeze(0), - None, - ) - output[start:stop].copy_(segment[0]) + if query.device.type != "mps": + segment = self.forward( + query[start:stop].unsqueeze(0), + key[start:stop].unsqueeze(0), + value[start:stop].unsqueeze(0), + None, + ) + 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 diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 07b69cba2..bccc9229f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/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 ( - x.dtype != torch.float32 - or layer.weight.dtype != torch.float32 - or (bias is not None and bias.dtype != torch.float32) - ): + if x.device.type == "mps": + if x.dtype == layer.weight.dtype and ( + bias is None or bias.dtype == x.dtype + ): + return F.linear(x, layer.weight, bias) return F.linear( x.to(torch.float32), layer.weight.to(torch.float32), diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 1c395f2c6..25d4cdac5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index aa903c041..949cc9ab9 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/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_name + component_starts_on_cpu = ( + 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( - local_torch_device, - component_starts_on_cpu=component_starts_on_cpu, - runtime_quant_config=quant_spec.runtime_quant_config, + checkpoint_load_device = ( + torch.device("cpu") + if cpu_offload_flag + 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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 4c6f842cf..411964068 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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_name + component_starts_on_cpu = ( + 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()) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 0190c6403..ee8c3e678 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py b/python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py index dd7de205e..877723188 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 6261417bc..b86440759 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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 ( - RealtimeSessionCache, -) +from sglang.multimodal_gen.runtime.realtime.session import 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 = torch.get_device_module().max_memory_reserved() - peak_allocated_bytes = torch.get_device_module().max_memory_allocated() + peak_reserved_bytes = final_snapshot.peak_reserved_mb * (1024**2) + peak_allocated_bytes = final_snapshot.peak_allocated_mb * (1024**2) 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 = peak_reserved_bytes / (1024**2) + output_batch.peak_memory_mb = capture_memory_snapshot().peak_reserved_mb def _forward_group(self, batch: list[Req]) -> OutputBatch: assert self.pipeline is not None diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py index e9e6ee082..9cccc78bc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index adeae8a19..19c566c2d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/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.prefetch_size = min(max(1, prefetch_size), self.num_layers) + self._synchronous_mps = current_platform.is_mps() + # 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( - current_platform.device_type, torch.get_device_module().current_device() + self.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 - event = torch.get_device_module().Event() - event.record(self.copy_stream) - self._prefetch_events[layer_idx] = event + if self.copy_stream is not None: + # record after all copies so the consumer waits for every weight copy + event = torch.get_device_module().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", diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 392107ed8..4079ad8e8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index c3ff01c4c..c5c9a694d 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/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 " diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py index 80026a7e6..a270a5cd9 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py +++ b/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] diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py index 35daf9015..39d37376e 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py +++ b/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 ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index abbd6f6a4..d9f30791b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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(), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py index 155d5ff4a..22aa3a4ce 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py index eeba8b353..bd8d0e7d0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py index 130c8fad4..33d80ad94 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py +++ b/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( - device_type=audio_latent.device.type, - dtype=audio_vae_dtype, - enabled=audio_autocast_enabled, - ): + autocast_context = ( + torch.autocast( + device_type="cuda", + 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( - device_type=visual_latent.device.type, - dtype=video_vae_dtype, - enabled=visual_autocast_enabled, - ): + autocast_context = ( + torch.autocast( + device_type="cuda", + 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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py index 69cd37352..381d257e4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py +++ b/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(): - raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA") - device = torch.device("cuda") + if not (current_platform.is_cuda() or current_platform.is_mps()): + raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA or MPS") + 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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py index 025e8d4d9..7ac916e21 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/platforms/mps.py b/python/sglang/multimodal_gen/runtime/platforms/mps.py index ec1e5ded9..d4547407d 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/mps.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 7b709cb7e..8e2865b88 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/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 components" + "MPS supports only resident or layerwise-offload component " + "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 diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index cd74836bd..e408c4ad7 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/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}" diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 6346e1583..0c5ebaf3b 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index 978df8ba6..a894d3c50 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/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)