Merge commit '6fe1ff7709c72c08bbfc33a7e14b73ec94135549' into codex-issue-7035-h100-determinism

Signed-off-by: david6666666 <530634352@qq.com>

vllm_omni/diffusion/models/minimax_h3/fasth3.py

@@ -102,6 +102,7 @@ _MODEL_LEVEL_TARGETS = {
     "time_embedder.linear_2": "time_embedder.proj_out",
     "norm_out.linear": "final_layer.adaln_proj.linear",
     "norm_out.norm": "final_layer.norm",
+    "token_refiner.final_norm": "token_refiner.final_norm",
 }
 
 # Per-block adapter suffix -> (native suffix, how a delta enters the native
@@ -115,6 +116,8 @@ _BLOCK_TARGETS = {
     "attn.to_v": ("attn.qkv_proj", _QKV_V),
     "attn.to_out.0": ("attn.out_proj", _PLAIN),
     "attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN),
+    "attn.norm_q": ("attn.q_norm", _PLAIN),
+    "attn.norm_k": ("attn.k_norm", _PLAIN),
     "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
     "ff.net.2": ("mlp.fc2", _PLAIN),
     "adaln_proj.linear": ("adaln_proj.linear", _PLAIN),

vllm_omni/diffusion/models/minimax_h3/fasth3_checkpoint.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Sampling contract for FastVideo FastH3 full checkpoints."""
+
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2"
+FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0))
+
+
+@dataclass(frozen=True)
+class FastH3CheckpointSpec:
+    """The full V2 release has its own schedule and trained attention policy.
+
+    This is independent of ``FastH3WeightFusion``: the released weights already
+    contain the entire student, including its learned compression gates.
+    """
+
+    vsa_sparsity: float = 0.8
+
+    @classmethod
+    def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec:
+        """Validate the release's own fastvideo_inference.json."""
+        required = {
+            "model_id": FASTH3_V2_MODEL_ID,
+            "schema_version": "fasth3-inference-contract-v1",
+            "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
+            "transformer_forwards": 8,
+            "video_scheduler_shift": 10.0,
+            "audio_scheduler_shift": 3.0,
+            "guidance_scale": 1.0,
+            "attention_backend": "VIDEO_SPARSE_ATTN_H3",
+            "vsa_sparsity": 0.8,
+            "vsa_tile_size": 64,
+            "task": "t2av",
+        }
+        for key, expected in required.items():
+            if metadata.get(key) != expected:
+                raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}")
+        return cls()
+
+    def release_metadata(self) -> dict[str, Any]:
+        """Express the sampling policy in the existing H3 pipeline schema."""
+        return {
+            "partition": "fl2va",
+            "tasks": ["t2va"],
+            "sigma_shift_scales": {"video": 10.0, "audio": 3.0},
+            "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
+        }
+
+    def resolve_native_vaes(self, model_root: Path) -> Path:
+        """Reuse the frozen base VAEs with Omni's native tiled/parallel runtime.
+
+        Only VAE components are fetched; the student and text encoder come
+        directly from the FastVideo release. The release pins the base revision.
+        """
+        provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8"))
+        base = provenance.get("base_model", "")
+        prefix = f"hf://{FASTH3_BASE_MODEL}@"
+        if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]:
+            raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision")
+        return (
+            Path(
+                download_weights_from_hf_specific(
+                    model_name_or_path=FASTH3_BASE_MODEL,
+                    cache_dir=None,
+                    allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
+                    revision=base[len(prefix) :],
+                    require_all=True,
+                )
+            )
+            / "FL2VA"
+        )
+
+    def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None:
+        if partition != "fl2va":
+            raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)")
+        if getattr(od_config, "lora_path", None):
+            raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported")
+        if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA":
+            raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA")
+        attention_config = getattr(od_config, "diffusion_attention_config", None)
+        per_role = getattr(attention_config, "per_role", None) or {}
+        spec = per_role.get("self") or getattr(attention_config, "default", None)
+        if getattr(spec, "fastvideo_vsa_topk", None) is not None:
+            raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override")
+        parallel = getattr(od_config, "parallel_config", None)
+        if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")):
+            raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism")
+
+    def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None:
+        if step_execution:
+            # StepScheduler admits requests before the pipeline hook runs and
+            # derives its lifetime from these fields.  The pinned schedule must
+            # therefore be explicit and cannot be overridden by custom arrays.
+            if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
+                raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8")
+            if sampling.timesteps is not None or sampling.sigmas is not None:
+                raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas")
+        if sampling.lora_request is not None:
+            raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters")
+        steps = sampling.num_inference_steps
+        if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
+            raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted")
+        extra = sampling.extra_args or {}
+        for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)):
+            try:
+                value = float(extra.get(key, expected))
+            except (TypeError, ValueError) as exc:
+                raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc
+            if not math.isclose(value, expected):
+                raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}")
+        if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0:
+            raise OmniClientError("FastH3 V2 requires guidance_scale=1")

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -55,6 +55,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.platforms import current_omni_platform
 
+from .fasth3 import _resolve_native_target
+
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
         QuantizationConfig,
@@ -118,6 +120,18 @@ class MiniMaxH3DiTArchConfig:
 
     @classmethod
     def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig:
+        # The modular Diffusers checkpoint uses these spellings for the same
+        # architecture. Normalize before constructing the existing native DiT.
+        aliases = {
+            "num_refiner_layers": "token_refiner_num_layers",
+            "ffn_dim": "ffn_hidden_size",
+            "in_channels": "latents_dim",
+            "audio_in_channels": "audio_latents_dim",
+            "freq_dim": "timestep_input_dim",
+            "time_embed_hidden_dim": "time_embed_hidden_size",
+            "rope_freq_dim": "rope_inv_freq_len",
+        }
+        config = {aliases.get(name, name): value for name, value in config.items()}
         fields = cls.__dataclass_fields__
         values = {name: config[name] for name in fields if name in config}
         if "patch_size" in values:
@@ -394,6 +408,9 @@ def _sdpa_varlen_attention(
 
 
 class MiniMaxH3Attention(nn.Module):
+    # Full sparse checkpoints pin a ratio; legacy adapters use backend top-k.
+    vsa_sparsity: float | None = None
+
     def __init__(
         self,
         arch: MiniMaxH3DiTArchConfig,
@@ -588,6 +605,7 @@ class MiniMaxH3Attention(nn.Module):
                 # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other
                 # backend/path.
                 "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE,
+                **({"vsa_h3_sparsity": self.vsa_sparsity} if self.vsa_sparsity is not None else {}),
                 # Present only for a VSA artifact; the VSA backend reads it as
                 # the learned compression gate and every other backend ignores it.
                 **({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}),
@@ -1127,10 +1145,23 @@ class MiniMaxH3DiTModel(nn.Module):
         self,
         od_config: OmniDiffusionConfig,
         quant_config: QuantizationConfig | None = None,
+        *,
+        diffusers_weights: bool | None = None,
     ) -> None:
         super().__init__()
         tf_config = od_config.tf_model_config
         config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config)
+        # The native MiniMax-H3 Hub snapshot advertises the Diffusers
+        # transformer class in its root config, while its FL2VA/Ref2VA
+        # components still contain native weights.  The pipeline has already
+        # resolved the actual source format, so let it override the
+        # class-name heuristic.  Keep the heuristic for standalone callers.
+        self._diffusers_weights = (
+            config_mapping.get("_class_name") == "MiniMaxH3Transformer3DModel"
+            if diffusers_weights is None
+            else diffusers_weights
+        )
+        self._rope_theta = float(config_mapping.get("rope_theta", 10000.0))
         arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping)
         self.arch = arch
         self.od_config = od_config
@@ -1211,7 +1242,7 @@ class MiniMaxH3DiTModel(nn.Module):
         )
         self._mark_missing_params_required()
 
-    def enable_vsa_gates(self) -> None:
+    def enable_vsa_gates(self, *, sparsity: float | None = None) -> None:
         """Give every DiT block's attention a VSA compression gate.
 
         A FastH3 VSA artifact assigns these projections rather than adding to
@@ -1223,6 +1254,9 @@ class MiniMaxH3DiTModel(nn.Module):
             return
         for block in self.blocks:
             block.attn.enable_vsa_gate()
+            block.attn.vsa_sparsity = sparsity
+            if sparsity is not None:
+                block.attn.to_gate_compress.weight.missing_param_init = "error"
         self.vsa_gates_enabled = True
 
     def _mark_missing_params_required(self) -> None:
@@ -1293,17 +1327,37 @@ class MiniMaxH3DiTModel(nn.Module):
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
     ) -> set[str]:
-        """Load exact H3 checkpoint names with logical TP-aware loaders."""
+        """Load native or Diffusers H3 weights with the existing TP-aware loaders."""
         params = dict(self.named_parameters())
         params.update(dict(self.named_buffers()))
         loaded: set[str] = set()
+        diffusers_weights = getattr(self, "_diffusers_weights", False)
+        qkv_parts: dict[str, set[str]] = {}
+        source_names: set[str] = set()
         for name, loaded_weight in weights:
+            layout = "plain"
+            if diffusers_weights:
+                if name in source_names:
+                    raise ValueError(f"duplicate Diffusers H3 weight: {name}")
+                source_names.add(name)
+                module, _, kind = name.rpartition(".")
+                target = _resolve_native_target(module)
+                if target is None or kind not in {"weight", "bias"}:
+                    raise ValueError(f"unsupported Diffusers H3 weight: {name}")
+                name, layout = f"{target[0]}.{kind}", target[1]
             param = params.get(name)
             if param is None:
+                if diffusers_weights:
+                    raise ValueError(f"Diffusers H3 weight has no model parameter: {name}")
                 logger.warning("Skipping MiniMax H3 weight not present in model: %s", name)
                 continue
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
-            if name.endswith(".attn.qkv_proj.weight"):
+            if layout in {"q", "k", "v"}:
+                # vLLM can load each projection directly into its packed QKV
+                # parameter, including TP slicing and online quantization.
+                weight_loader(param, loaded_weight, layout)
+                qkv_parts.setdefault(name, set()).add(layout)
+            elif name.endswith(".attn.qkv_proj.weight"):
                 # Transform checkpoint layout before entering vLLM's loader so
                 # online FP8 can keep ``online_process_loader`` outermost.
                 loaded_weight = _reorder_grouped_qkv_to_qkv(
@@ -1319,12 +1373,26 @@ class MiniMaxH3DiTModel(nn.Module):
                         "MiniMax H3 fc1 checkpoint rows must split evenly into "
                         f"gate/up matrices, got {tuple(loaded_weight.shape)}"
                     )
-                gate, up = loaded_weight.chunk(2, dim=0)
+                first, second = loaded_weight.chunk(2, dim=0)
+                gate, up = (second, first) if layout == "swap_halves" else (first, second)
                 weight_loader(param, gate, 0)
                 weight_loader(param, up, 1)
             else:
                 weight_loader(param, loaded_weight)
             loaded.add(name)
+        if diffusers_weights:
+            for name, parts in qkv_parts.items():
+                if parts != {"q", "k", "v"}:
+                    raise ValueError(f"incomplete Diffusers H3 QKV group {name}: {sorted(parts)}")
+            # Diffusers reconstructs RoPE from config instead of storing this
+            # native checkpoint buffer. Compute on CPU for identical values.
+            freq_dim = self.arch.rope_inv_freq_len
+            rope = 1.0 / (
+                self._rope_theta
+                ** (torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device="cpu") / (2 * freq_dim))
+            )
+            default_weight_loader(params["rope.inv_freq"], rope)
+            loaded.add("rope.inv_freq")
         return loaded
 
     @staticmethod

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -67,6 +67,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
 from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
+    is_minimax_h3_modular,
     resolve_minimax_h3_partition,
 )
 from vllm_omni.model_executor.models.minimax_h3.conditioning import (
@@ -106,6 +107,7 @@ from .denoise_loop import (
 )
 from .encoder import MiniMaxH3Qwen3VLEncoder
 from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion
+from .fasth3_checkpoint import FastH3CheckpointSpec
 from .lora import TurboSpec, load_minimax_h3_turbo_lora
 from .minimax_h3_transformer import (
     MiniMaxH3Attention,
@@ -213,7 +215,11 @@ def _resolve_minimax_h3_model_root(
         if path.name in {"FL2VA", "Ref2VA"}:
             return path.parent
         return path
-    if load_text_encoder:
+    if is_minimax_h3_modular(model, revision):
+        allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"]
+        if load_text_encoder:
+            allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"]
+    elif load_text_encoder:
         allow_patterns = (
             MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
         )
@@ -314,6 +320,8 @@ def resolve_minimax_h3_diffusion_model_path(
         partition,
         load_text_encoder=False,
     )
+    if is_minimax_h3_modular(str(model_root), revision):
+        return str(model_root)
     if partition == "combined":
         return str(model_root)
     subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
@@ -575,6 +583,7 @@ class MiniMaxH3Pipeline(
     _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
     # Set from --lora-path during construction; absent means no FastH3 adapter.
     _fasth3: FastH3WeightFusion | None = None
+    _fasth3_checkpoint: FastH3CheckpointSpec | None = None
 
     def _load_diffusion_lora_adapter(
         self,
@@ -800,10 +809,13 @@ class MiniMaxH3Pipeline(
             self._PROFILER_TARGETS.remove("encode_prompt")
         if not self.load_vae_encoder:
             self._PROFILER_TARGETS.remove("_encode_local_media")
+        modular = is_minimax_h3_modular(str(od_config.model), od_config.revision)
         self.partition = _minimax_h3_partition_for_task(
             getattr(od_config, "task_type", None),
             str(od_config.model),
         )
+        if modular and str(od_config.task_type or "auto").lower() == "auto":
+            self.partition = "fl2va"
         self._turbo_lora_specs: dict[int, TurboSpec] = {}
         self._native_lora_adapter_ids: set[int] = set()
         self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {}
@@ -813,9 +825,19 @@ class MiniMaxH3Pipeline(
             self.partition,
             load_text_encoder=self.load_text_encoder,
         )
-        model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
-        model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
-        release = model_index.get("_minimax_h3") or {}
+        if modular:
+            model_path = model_root
+            self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(
+                json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8"))
+            )
+            self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config)
+            release = self._fasth3_checkpoint.release_metadata()
+            vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root)
+        else:
+            model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
+            model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
+            release = model_index.get("_minimax_h3") or {}
+            vae_model_path = model_path
         partition = str(release.get("partition", "")).lower()
         expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va"
         if partition != expected_partition:
@@ -876,11 +898,18 @@ class MiniMaxH3Pipeline(
         self.transformer = MiniMaxH3DiTModel(
             od_config,
             quant_config=transformer_quant_config,
+            diffusers_weights=modular,
         )
+        if self._fasth3_checkpoint is not None:
+            self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity)
+            logger.info(
+                "FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64"
+            )
         if ref2va_model_path is not None:
             self.transformers_ref = MiniMaxH3DiTModel(
                 od_config,
                 quant_config=transformer_quant_config,
+                diffusers_weights=modular,
             )
 
         self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
@@ -963,14 +992,14 @@ class MiniMaxH3Pipeline(
         # so VAEs stay resident for new configurations.
         component_load_device = torch.device("cpu") if legacy_manual_components else self.device
         self.video_vae = MiniMaxH3VideoVAE(
-            os.path.join(model_path, "video_vae"),
+            os.path.join(vae_model_path, "video_vae"),
             device=self.device,
             load_device=component_load_device,
             decode_only=not self.load_vae_encoder,
             trust_remote_code=od_config.trust_remote_code,
         )
         self.audio_vae = MiniMaxH3AudioVAE(
-            os.path.join(model_path, "audio_vae"),
+            os.path.join(vae_model_path, "audio_vae"),
             device=self.device,
             load_device=component_load_device,
             decode_only=not self.load_vae_encoder,
@@ -1046,6 +1075,12 @@ class MiniMaxH3Pipeline(
             # load_weights only warns on a parameter the model does not have, so
             # close the adapter against what the DiT actually consumed.
             self._fasth3.validate_fully_applied(transformer_loaded)
+        if self._fasth3_checkpoint is not None:
+            required_gates = {
+                f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers)
+            }
+            if missing_gates := required_gates - transformer_loaded:
+                raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}")
         return loaded_with_prefix
 
     @property
@@ -2056,6 +2091,11 @@ class MiniMaxH3Pipeline(
             self._validate_turbo_sampling(sampling, turbo_spec)
         if has_native_lora:
             self._validate_native_sampling(sampling, task=task)
+        if self._fasth3_checkpoint is not None:
+            self._fasth3_checkpoint.check_request(
+                sampling,
+                step_execution=bool(getattr(self.od_config, "step_execution", False)),
+            )
         if self._fasth3 is not None:
             self._fasth3.check_request(
                 sampling,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
     return {}
 
 
+def resolve_wan_transformer_quant_config(
+    config: dict, quant_config: QuantizationConfig | None, component: str
+) -> QuantizationConfig | None:
+    """Resolve the expert before applying its checkpoint's storage contract."""
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+    # Wan experts are siblings: "transformer_2" must not inherit "transformer"
+    # through the generic layer-prefix resolver. Select the whole expert name.
+    component_quant_config = (
+        quant_config.component_configs.get(component, quant_config.default_config)
+        if isinstance(quant_config, ComponentQuantizationConfig)
+        else quant_config
+    )
+    quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None
+    resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config"))
+    if quantization_disabled and resolved_quant_config is not None:
+        raise ValueError(
+            f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. "
+            "Use a BF16 checkpoint for this component or enable its matching quantization method."
+        )
+    return resolved_quant_config
+
+
 def create_transformer_from_config(
-    config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
+    config: dict,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+    *,
+    component: str = "transformer",
 ) -> WanTransformer3DModel:
     """Create WanTransformer3DModel from config dict."""
     kwargs: dict = {}
@@ -216,11 +244,7 @@ def create_transformer_from_config(
     if "pos_embed_seq_len" in config:
         kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
 
-    if "quantization_config" in config:
-        from vllm_omni.quantization.factory import resolve_quant_config_from_disk
-
-        quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
-
+    quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
     if quant_config is not None:
         kwargs["quant_config"] = quant_config
     if prefix:
@@ -460,13 +484,13 @@ class Wan22Pipeline(
         # Initialize transformers with correct config (weights loaded via load_weights)
         if load_transformer:
             transformer_config = load_transformer_config(model, "transformer", local_files_only)
-            self.transformer = self._create_transformer(transformer_config)
+            self.transformer = self._create_transformer(transformer_config, component="transformer")
         else:
             self.transformer = None
 
         if load_transformer_2:
             transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
-            self.transformer_2 = self._create_transformer(transformer_2_config)
+            self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
         else:
             self.transformer_2 = None
 
@@ -504,10 +528,17 @@ class Wan22Pipeline(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
-    def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+    def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
         """Create a transformer from a config dict. Respects od_config.quantization_config."""
         quant_config = getattr(self.od_config, "quantization_config", None)
-        return create_transformer_from_config(config, quant_config=quant_config)
+        # Startup metadata describes the first expert, not a user policy for both.
+        if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+            quant_config = None
+        return create_transformer_from_config(
+            config,
+            quant_config=quant_config,
+            component=component,
+        )
 
     @property
     def guidance_scale(self):

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -304,25 +304,10 @@ class Wan22I2VPipeline(
         # Transformers (weights loaded via load_weights)
         # Load config from model directory or HF Hub to get correct in_channels for I2V models
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(
-            transformer_config,
-            quant_config=od_config.quantization_config,
-        )
+        self.transformer = self._create_transformer(transformer_config, component="transformer")
         if self.has_transformer_2:
             transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
-            t2_quant = transformer_2_config.get("quantization_config")
-            if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
-                from vllm_omni.quantization.factory import build_quant_config
-
-                method = t2_quant["quant_method"]
-                kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
-                t2_quant = build_quant_config(method, **kwargs)
-            else:
-                t2_quant = None
-            self.transformer_2 = create_transformer_from_config(
-                transformer_2_config,
-                quant_config=t2_quant,
-            )
+            self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2")
         else:
             self.transformer_2 = None
 
@@ -462,10 +447,13 @@ class Wan22I2VPipeline(
         image_embeds = self.image_encoder(pixel_values, output_hidden_states=True)
         return image_embeds.hidden_states[-2]
 
-    def _create_transformer(self, config: dict) -> WanTransformer3DModel:
+    def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel:
         """Create a transformer from a config dict. Respects od_config.quantization_config."""
         quant_config = getattr(self.od_config, "quantization_config", None)
-        return create_transformer_from_config(config, quant_config=quant_config)
+        # Startup metadata describes the first expert, not a user policy for both.
+        if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+            quant_config = None
+        return create_transformer_from_config(config, quant_config=quant_config, component=component)
 
     def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         sampling_params_list = req.sampling_params_list

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     Wan22Pipeline,
     load_wan_weights_with_optional_gate,
+    resolve_wan_transformer_quant_config,
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -60,6 +61,8 @@ def create_vace_transformer_from_config(
     config: dict,
     quant_config: QuantizationConfig | None = None,
     prefix: str = "",
+    *,
+    component: str = "transformer",
 ) -> WanVACETransformer3DModel:
     """Create WanVACETransformer3DModel from config dict."""
     kwargs = {}
@@ -97,6 +100,7 @@ def create_vace_transformer_from_config(
         kwargs["vace_layers"] = config["vace_layers"]
     if "vace_in_channels" in config:
         kwargs["vace_in_channels"] = config["vace_in_channels"]
+    quant_config = resolve_wan_transformer_quant_config(config, quant_config, component)
     if quant_config is not None:
         kwargs["quant_config"] = quant_config
     if prefix:
@@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
 
         super().__init__(od_config=od_config, prefix=prefix)
 
-    def _create_transformer(self, config: dict) -> WanVACETransformer3DModel:
+    def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel:
         """Build VACE transformer. Respects od_config.quantization_config."""
         quant_config = getattr(self.od_config, "quantization_config", None)
-        return create_vace_transformer_from_config(config, quant_config=quant_config)
+        # Startup metadata describes the first expert, not a user policy for both.
+        if getattr(self.od_config, "quantization_config_is_auto_detected", False):
+            quant_config = None
+        return create_vace_transformer_from_config(config, quant_config=quant_config, component=component)
 
     def diffuse(
         self,

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -18,7 +18,7 @@ from vllm.distributed import (
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.conv import Conv3dLayer
-from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.utils import (
@@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
 from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod
 
 logger = init_logger(__name__)
 
@@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module):
         # ROPE helper
         self._cached_rope_emb = None
         self._cached_rope_resolution = None
+        self._validate_w4a8_fallback_layers(quant_config)
+
+    def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None:
+        """Reject checkpoint aliases/typos before loading weights or running inference."""
+        requested = tuple(getattr(quant_config, "w4a8_fallback_layers", []))
+        self._w4a8_fallback_layers = requested
+        self._local_w4a8_fallback_layers: tuple[str, ...] = ()
+        if not requested:
+            return
+        modules = dict(self.named_modules())
+        missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)]
+        local_layers = [
+            name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes)
+        ]
+        self._local_w4a8_fallback_layers = tuple(local_layers)
+        invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)]
+        if invalid:
+            raise ValueError(
+                "w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; "
+                f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, "
+                "not checkpoint attn1.to_q/to_k/to_v names."
+            )
+
+    def _log_w4a8_fallback_load_summary(self, component_name: str) -> None:
+        """Log post-load transforms/cache readiness, not numerical acceptance."""
+        requested = self._w4a8_fallback_layers
+        if not requested:
+            return
+        modules = dict(self.named_modules())
+        local_layers = self._local_w4a8_fallback_layers
+        selected_layers = [
+            name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False)
+        ]
+        bf16_overrides = [name for name in local_layers if name not in selected_layers]
+        processed = [
+            name
+            for name in selected_layers
+            if getattr(modules[name], "_already_called_process_weights_after_loading", False)
+        ]
+        if not selected_layers:
+            processing_state = "not-required"
+        elif len(processed) == len(selected_layers):
+            processing_state = "ready"
+        else:
+            processing_state = "not-ready"
+
+        def qualified(names: Iterable[str]) -> list[str]:
+            return [f"{component_name}.{name}" for name in names]
+
+        logger.info(
+            "Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; "
+            "weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d",
+            component_name,
+            len(selected_layers),
+            qualified(selected_layers),
+            len(bf16_overrides),
+            qualified(bf16_overrides),
+            processing_state,
+            len(processed),
+            len(selected_layers),
+            len(requested) - len(local_layers),
+        )
 
     @property
     def dtype(self) -> torch.dtype:
@@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module):
         }
 
         params_dict = dict(self.named_parameters())
+        single_scale_prefixes = {
+            name
+            for name, module in self.named_modules()
+            if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod)
+        }
+
+        def reject_dualscale_tensor(name: str) -> None:
+            if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes:
+                raise ValueError(
+                    f"Single-scale mxfp4 cannot load DualScale tensor {name}; "
+                    "changing quant_method does not convert a checkpoint."
+                )
+
         loaded_params: set[str] = set()
+        loaded_qkv_shards: dict[str, set[str]] = {}
+        qkv_smooth_scales: dict[str, torch.Tensor] = {}
 
         for name, loaded_weight in weights:
             name = weight_name_remapping.get(name, name)
@@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module):
             # Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall
             # through to the else branch and are loaded directly.
             for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in original_name:
+                if f"{weight_name}." not in original_name:
                     continue
-                lookup_name = original_name.replace(weight_name, param_name)
+                lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1)
+                reject_dualscale_tensor(lookup_name)
                 # Skip weights that belong to PP stages other than this one
                 if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
                     break
                 param = params_dict[lookup_name]
+                if lookup_name.endswith(".mul_scale"):
+                    previous_scale = qkv_smooth_scales.get(lookup_name)
+                    if previous_scale is not None and not torch.equal(previous_scale, loaded_weight):
+                        raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}")
+                    if previous_scale is None:
+                        qkv_smooth_scales[lookup_name] = loaded_weight.clone()
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(original_name)
+                if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False):
+                    # Shared input-channel Smooth is a complete tensor even
+                    # when supplied under a single Q/K/V source name.
+                    loaded_params.add(lookup_name)
+                else:
+                    shards = loaded_qkv_shards.setdefault(lookup_name, set())
+                    shards.add(shard_id)
+                    if shards == {"q", "k", "v"}:
+                        loaded_params.add(lookup_name)
                 break
             else:
                 # diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
@@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module):
                 if is_pp_missing_parameter(lookup_name, self):
                     continue
 
+                reject_dualscale_tensor(lookup_name)
                 if lookup_name not in params_dict:
                     logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
                     continue
@@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module):
 
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
-
-            loaded_params.add(original_name)
-            loaded_params.add(lookup_name)
+                loaded_params.update((original_name, lookup_name))
 
         return loaded_params