[Model] Add native MAGI-2 Preview diffusion support (#5918)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Loader-side adapters for proven direct-checkpoint mmap layouts."""
 
 from __future__ import annotations
@@ -83,8 +83,30 @@ class _Cosmos3DirectMmapAdapter:
         return DirectMmapTensorPolicy(allow_custom_loader=True)
 
 
+class _Magi2DirectMmapAdapter:
+    """Apply MAGI-2's rank-local MoE-head slices during bounded staging."""
+
+    def policy_for(
+        self,
+        runtime_name: str,
+        target: torch.Tensor,
+    ) -> DirectMmapTensorPolicy:
+        del runtime_name
+        transform = getattr(target, "mmap_weight_transform", None)
+        return DirectMmapTensorPolicy(
+            transform=transform if callable(transform) else None,
+        )
+
+
 def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None:
     """Return a loader-owned adapter without requiring a pipeline flag."""
+    from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
+        Magi2PreviewTransformer,
+    )
+
+    if any(isinstance(module, Magi2PreviewTransformer) for module in pipeline.modules()):
+        return _Magi2DirectMmapAdapter()
+
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
         MiniMaxH3DiTModel,
     )

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -270,15 +270,48 @@ def _validate_source_metadata(
                 source = handle.get_slice(binding.checkpoint_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
-                if source_shape != tuple(target.shape):
-                    raise _PlanIncompatibleError(
-                        f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}"
-                    )
-                if source_dtype is None or source_dtype != target.dtype:
+                if source_dtype is None:
                     raise _PlanIncompatibleError(
                         f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}"
                     )
 
+                runtime_shape = source_shape
+                runtime_dtype = source_dtype
+                if binding.transform is not None:
+                    # Adapter transforms can map full checkpoint tensors to a
+                    # rank-local runtime shape. Validate that contract without
+                    # materializing the source tensor during preflight.
+                    try:
+                        transformed = binding.transform(
+                            torch.empty(
+                                source_shape,
+                                dtype=source_dtype,
+                                device="meta",
+                            )
+                        )
+                    except (NotImplementedError, RuntimeError, TypeError, ValueError) as exc:
+                        raise _PlanIncompatibleError(
+                            f"checkpoint transform for {runtime_name!r} cannot be validated on tensor metadata: {exc}"
+                        ) from exc
+                    if not isinstance(transformed, torch.Tensor):
+                        raise _PlanIncompatibleError(
+                            f"checkpoint transform for {runtime_name!r} returned {type(transformed).__name__}, "
+                            "expected torch.Tensor"
+                        )
+                    runtime_shape = tuple(transformed.shape)
+                    runtime_dtype = transformed.dtype
+
+                if runtime_shape != tuple(target.shape):
+                    raise _PlanIncompatibleError(
+                        f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, "
+                        f"transformed={runtime_shape}, runtime={tuple(target.shape)}"
+                    )
+                if runtime_dtype != target.dtype:
+                    raise _PlanIncompatibleError(
+                        f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, "
+                        f"transformed={runtime_dtype}, runtime={target.dtype}"
+                    )
+
 
 def build_checkpoint_mmap_plan(
     pipeline: nn.Module,