[diffusion][CI]: Add individual component accuracy CI for diffusion models (#18709)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

.github/workflows/pr-test-multimodal-gen.yml

@@ -156,6 +156,108 @@ jobs:
         with:
           artifact-suffix: ${{ matrix.part }}
 
+  multimodal-gen-component-accuracy-1-gpu:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-component-accuracy-1-gpu') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-stage-health
+
+      - uses: ./.github/actions/check-maintenance
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run diffusion component accuracy tests (1-GPU)
+        timeout-minutes: 240
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+        run: |
+          cd python
+          if [ "${{ inputs.continue_on_error }}" = "true" ]; then
+            exit_code=0
+            python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py || exit_code=$?
+            python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py || exit_code=$?
+            exit $exit_code
+          fi
+          python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py
+          python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: always()
+
+  multimodal-gen-component-accuracy-2-gpu:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-component-accuracy-2-gpu') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: 2-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-stage-health
+
+      - uses: ./.github/actions/check-maintenance
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run diffusion component accuracy tests (2-GPU)
+        timeout-minutes: 240
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+        run: |
+          cd python
+          if [ "${{ inputs.continue_on_error }}" = "true" ]; then
+            exit_code=0
+            torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py || exit_code=$?
+            torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py || exit_code=$?
+            exit $exit_code
+          fi
+          torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py
+          torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: always()
+
   multimodal-gen-unit-test:
     if: |
       (inputs.target_stage == 'multimodal-gen-unit-test') ||

.github/workflows/pr-test.yml

@@ -890,6 +890,8 @@ jobs:
       (
         inputs.target_stage == 'multimodal-gen-test-1-gpu' ||
         inputs.target_stage == 'multimodal-gen-test-2-gpu' ||
+        inputs.target_stage == 'multimodal-gen-component-accuracy-1-gpu' ||
+        inputs.target_stage == 'multimodal-gen-component-accuracy-2-gpu' ||
         inputs.target_stage == 'multimodal-gen-unit-test' ||
         (
           !inputs.target_stage &&

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

@@ -156,7 +156,9 @@ class TextEncoderLoader(ComponentLoader):
         return hf_folder, hf_weights_files, use_safetensors
 
     def _get_weights_iterator(
-        self, source: "Source", to_cpu: bool
+        self,
+        source: "Source",
+        to_cpu: bool,
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
         """get an iterator for the model weights based on the load format."""
         hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
@@ -166,7 +168,8 @@ class TextEncoderLoader(ComponentLoader):
         )
         if use_safetensors:
             weights_iterator = safetensors_weights_iterator(
-                hf_weights_files, to_cpu=to_cpu
+                hf_weights_files,
+                to_cpu=to_cpu,
             )
         else:
             weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
@@ -186,17 +189,27 @@ class TextEncoderLoader(ComponentLoader):
             fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
             allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
         )
-        yield from self._get_weights_iterator(primary_weights, to_cpu)
+        yield from self._get_weights_iterator(
+            primary_weights,
+            to_cpu,
+        )
 
         secondary_weights = cast(
             Iterable[TextEncoderLoader.Source],
             getattr(model, "secondary_weights", ()),
         )
         for source in secondary_weights:
-            yield from self._get_weights_iterator(source, to_cpu)
+            yield from self._get_weights_iterator(
+                source,
+                to_cpu,
+            )
 
     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 | None = None,
     ):
         """Load the text encoders based on the model path, and inference args."""
         diffusers_pretrained_config = get_config(
@@ -227,6 +240,7 @@ class TextEncoderLoader(ComponentLoader):
             encoder_config,
             server_args,
             encoder_dtype,
+            cpu_offload_flag=cpu_offload_flag,
         )
 
     def load_model(
@@ -240,7 +254,10 @@ class TextEncoderLoader(ComponentLoader):
         # Determine CPU offload behavior and target device
 
         local_torch_device = get_local_torch_device()
-        should_offload = self.should_offload(server_args, model_config)
+        fsdp_cpu_offload = self.should_offload(server_args, model_config)
+        should_offload = (
+            cpu_offload_flag if cpu_offload_flag is not None else fsdp_cpu_offload
+        )
 
         if should_offload and not current_platform.is_mps():
             model_device = torch.device("cpu")
@@ -263,12 +280,12 @@ class TextEncoderLoader(ComponentLoader):
 
             weights_to_load = {name for name, _ in model.named_parameters()}
             loaded_weights = model.load_weights(
-                self._get_all_weights(model, model_path, to_cpu=should_offload)
+                self._get_all_weights(
+                    model,
+                    model_path,
+                    to_cpu=should_offload,
+                )
             )
-
-            # Explicitly move model to target device after loading weights
-            if not should_offload:
-                model = model.to(local_torch_device)
 
             if should_offload:
                 # Disable FSDP for MPS as it's not compatible
@@ -277,7 +294,7 @@ class TextEncoderLoader(ComponentLoader):
                         "Disabling FSDP sharding for MPS platform as it's not compatible"
                     )
                     model = model.to(local_torch_device)
-                else:
+                elif fsdp_cpu_offload:
                     mesh = init_device_mesh(
                         current_platform.device_type,
                         mesh_shape=(1, dist.get_world_size()),
@@ -292,6 +309,8 @@ class TextEncoderLoader(ComponentLoader):
                         or getattr(model, "_fsdp_shard_conditions", None),
                         pin_cpu_memory=server_args.pin_cpu_memory,
                     )
+                else:
+                    model = model.to("cpu")
             else:
                 model = model.to(local_torch_device)
             # We only enable strict check for non-quantized models

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

@@ -24,6 +24,7 @@ try:
 except ImportError:
     HAS_RUNAI_MODEL_STREAMER = False
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -135,13 +136,17 @@ def _validate_safetensors_file(file_path: str) -> bool:
 def safetensors_weights_iterator(
     hf_weights_files: list[str],
     to_cpu: bool = True,
-    use_runai_model_streamer: bool = HAS_RUNAI_MODEL_STREAMER,
+    use_runai_model_streamer: bool | None = None,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files."""
     enable_tqdm = (
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
     device = "cpu" if to_cpu else str(get_local_torch_device())
+    if use_runai_model_streamer is None:
+        use_runai_model_streamer = (
+            HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER
+        )
 
     # Validate files before loading
     corrupted_files = [

python/sglang/multimodal_gen/test/server/accuracy_config.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, Optional
+
+from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+
+
+class ComponentType(str, Enum):
+    VAE = "vae"
+    TRANSFORMER = "transformer"
+    TEXT_ENCODER = "text_encoder"
+
+
+@dataclass(frozen=True)
+class ComponentSkip:
+    reason: str
+
+
+DEFAULT_TIMESTEP = 500.0
+TIMESTEP_NORMALIZATION_FACTOR = 1000.0
+I2V_IMAGE_DIM = 1280
+I2V_TEXT_ENCODER_DIM = 5120
+
+DEFAULT_TEXT_ENCODER_VOCAB_SIZE = 32000
+TEXT_ENCODER_INPUT_SEED = 42
+TEXT_ENCODER_TOKEN_MIN = 100
+TEXT_ENCODER_TOKEN_MAX = 30000
+TEXT_ENCODER_TOKEN_LENGTH = 32
+
+# Default thresholds by component. Override per component/case if needed.
+DEFAULT_THRESHOLDS = {
+    ComponentType.VAE: 0.999,
+    ComponentType.TRANSFORMER: 0.995,
+    ComponentType.TEXT_ENCODER: 0.98,
+}
+
+# Optional per-case overrides: {case_id: {ComponentType: threshold}}
+CASE_THRESHOLDS: Dict[str, Dict[ComponentType, float]] = {
+    # Add overrides here when a specific model/component needs a different threshold.
+    "flux_2_image_t2i": {ComponentType.TRANSFORMER: 0.99},
+    "flux_2_image_t2i_layerwise_offload": {ComponentType.TRANSFORMER: 0.99},
+    "flux_2_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.99},
+    "flux_2_klein_ti2i_2_gpus": {ComponentType.TRANSFORMER: 0.975},
+    "flux_2_ti2i": {ComponentType.TRANSFORMER: 0.99},
+    "flux_2_t2i_customized_vae_path": {ComponentType.TRANSFORMER: 0.99},
+    "fast_hunyuan_video": {ComponentType.TRANSFORMER: 0.99},
+    "fsdp-inference": {ComponentType.TRANSFORMER: 0.9935},
+    "wan2_2_i2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99},
+    "wan2_2_t2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99},
+    "wan2_2_t2v_a14b_teacache_2gpu": {ComponentType.TRANSFORMER: 0.99},
+    "wan2_2_t2v_a14b_lora_2gpu": {ComponentType.TRANSFORMER: 0.99},
+    "zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.9935},
+    "zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.9935},
+}
+
+# Active skip policy. Keep this limited to cases with current, concrete evidence
+# of real divergence or unsupported reference loading in the harness.
+SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
+    "flux_image_t2i": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline despite 100% matched weights (CosSim ~0.47)"
+        )
+    },
+    "sana_image_t2i": {
+        ComponentType.VAE: ComponentSkip(
+            "HF AutoencoderDC checkpoint leaves required to_qkv_multiscale weights missing, so VAE transfer would compare against partially initialized reference weights"
+        )
+    },
+    "mova_360p_1gpu": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "HF reference transformer cannot be materialized from the video_dit repo layout"
+        )
+    },
+    "qwen_image_t2i_cache_dit_enabled": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by qwen_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by qwen_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by qwen_image_t2i for the same source component and topology"
+        ),
+    },
+    "flux_2_image_t2i_upscaling_4x": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+    },
+    "layerwise_offload": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+    },
+    "zimage_image_t2i_fp8": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+    },
+    "zimage_image_t2i_multi_lora": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
+        ),
+    },
+    "flux_2_ti2i": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+    },
+    "flux_2_t2i_customized_vae_path": {
+        ComponentType.VAE: ComponentSkip(
+            "Customized VAE override points to FLUX.2 Tiny AutoEncoder, but the HF reference loader does not yet materialize a trustworthy matching VAE baseline"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1.3b_teacache_enabled": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1.3b_frame_interp_2x": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1.3b_upscaling_4x": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "wan2_1_t2v_1_3b_lora_1gpu": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
+        ),
+    },
+    "flux_2_ti2i_multi_image_cache_dit": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
+        ),
+    },
+    "wan2_2_ti2v_5b": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "SGLang transformer loader rejects new parameters in HF checkpoint"
+        )
+    },
+    "fastwan2_2_ti2v_5b": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "SGLang transformer loader rejects new parameters in HF checkpoint"
+        )
+    },
+    "turbo_wan2_1_t2v_1.3b": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Weight transfer match ratio too low for reliable comparison"
+        )
+    },
+    "wan2_1_i2v_14b_480P_2gpu": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.71) after full weight transfer and matching output shape"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "wan2_1_i2v_14b_lora_2gpu": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_1_i2v_14b_720P_2gpu for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.68) after full weight transfer and matching output shape"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "wan2_1_i2v_14b_720P_2gpu": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.68) after full weight transfer and matching output shape"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "wan2_2_i2v_a14b_2gpu": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        )
+    },
+    "wan2_2_t2v_a14b_2gpu": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        )
+    },
+    "wan2_2_t2v_a14b_teacache_2gpu": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "wan2_2_t2v_a14b_lora_2gpu": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "wan2_1_t2v_14b_2gpu": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        )
+    },
+    "wan2_1_t2v_1.3b_cfg_parallel": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        )
+    },
+    "mova_360p_tp2": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "mova_360p_ring1_uly2": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "mova_360p_ring2_uly1": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
+        ),
+    },
+    "flux_image_t2i_2_gpus": {
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.47) after 100% matched weight transfer"
+        )
+    },
+    "zimage_image_t2i_2_gpus_non_square": {
+        ComponentType.VAE: ComponentSkip(
+            "Representative VAE accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "Representative transformer accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "Representative text encoder accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
+        ),
+    },
+    "flux_2_image_t2i_2_gpus": {
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "2-GPU FLUX.2 transformer diverges strongly from Diffusers baseline (CosSim ~0.54) despite full weight transfer"
+        )
+    },
+    "hunyuan3d_shape_gen": {
+        ComponentType.VAE: ComponentSkip(
+            "HF config cannot be parsed as valid JSON for component reference loading"
+        ),
+        ComponentType.TRANSFORMER: ComponentSkip(
+            "HF config cannot be parsed as valid JSON for component reference loading"
+        ),
+        ComponentType.TEXT_ENCODER: ComponentSkip(
+            "HF config cannot be parsed as valid JSON for component reference loading"
+        ),
+    },
+}
+
+# TODO: If a model needs extra compatibility logic, prefer adding a skip or an
+# explicit override here instead of adding more ad-hoc hacks in the engine.
+
+
+def get_threshold(case_id: str, component: ComponentType) -> float:
+    overrides = CASE_THRESHOLDS.get(case_id, {})
+    return overrides.get(component, DEFAULT_THRESHOLDS[component])
+
+
+def get_skip_reason(case: DiffusionTestCase, component: ComponentType) -> Optional[str]:
+    skip_entry = SKIP_COMPONENTS.get(case.id, {}).get(component)
+    if skip_entry is None:
+        return None
+    return skip_entry.reason
+
+
+def should_skip_component(case: DiffusionTestCase, component: ComponentType) -> bool:
+    return get_skip_reason(case, component) is not None

python/sglang/multimodal_gen/test/server/accuracy_hooks.py

@@ -0,0 +1,579 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Optional
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    DEFAULT_TIMESTEP,
+    I2V_IMAGE_DIM,
+    TIMESTEP_NORMALIZATION_FACTOR,
+    ComponentType,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    extract_output_tensor,
+    seed_and_broadcast,
+)
+
+Inputs = Dict[str, Any]
+BuildInputsFn = Callable[[Any, nn.Module, str, Optional[nn.Module]], Inputs]
+PrepareCallFn = Callable[[nn.Module, Inputs], "HookCall"]
+NormalizeFn = Callable[[Any], torch.Tensor]
+
+# These are harness defaults for synthetic accuracy inputs.
+# They are not checkpoint truth. We use them only when the model config or
+# forward signature does not expose a more specific shape or channel count.
+DEFAULT_TEXT_SEQ_LEN = 64
+DEFAULT_TOKEN_LAYOUT_SIZE = 32
+REDUCED_TOKEN_LAYOUT_SIZE = 16
+DEFAULT_VIDEO_FRAME_COUNT = 4
+DEFAULT_IMAGE_TOKEN_COUNT = 257
+ALIAS_ROTARY_TEXT_PAD_MULTIPLE = 32
+DEFAULT_TRANSFORMER_IN_CHANNELS = 16
+DEFAULT_TRANSFORMER_TEXT_CHANNELS = 4096
+DEFAULT_TRANSFORMER_POOLED_CHANNELS = 768
+DEFAULT_VAE_LATENT_CHANNELS = 16
+DEFAULT_VAE_LATENT_SPATIAL_SIZE = 32
+LARGE_CHANNEL_LAYOUT_THRESHOLD = 128
+
+
+@dataclass(frozen=True)
+class TransformerHookCompat:
+    normalize_reference_timestep: bool = False
+    negate_reference_output: bool = False
+    omit_reference_guidance: bool = False
+    use_2d_hidden_states: bool = False
+
+
+def _resolve_transformer_hook_compat(case: Any) -> TransformerHookCompat:
+    model_path = case.server_args.model_path.lower()
+    if "z-image" in model_path:
+        return TransformerHookCompat(
+            normalize_reference_timestep=True,
+            negate_reference_output=True,
+        )
+    if "qwen" in model_path:
+        return TransformerHookCompat(
+            normalize_reference_timestep=True,
+            omit_reference_guidance=True,
+        )
+    if "sana" in model_path:
+        return TransformerHookCompat(
+            omit_reference_guidance=True,
+            use_2d_hidden_states=True,
+        )
+    if "flux" in model_path:
+        return TransformerHookCompat(normalize_reference_timestep=True)
+    return TransformerHookCompat()
+
+
+@dataclass
+class HookCall:
+    module: nn.Module
+    args: tuple[Any, ...] = ()
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+    negate_output: bool = False
+
+
+@dataclass(frozen=True)
+class NativeHookProfile:
+    build_inputs: BuildInputsFn
+    prepare_sglang_call: PrepareCallFn
+    prepare_reference_call: PrepareCallFn
+    normalize_sglang_output: NormalizeFn = extract_output_tensor
+    normalize_reference_output: NormalizeFn = extract_output_tensor
+
+
+class _DeterministicRNG:
+    def __init__(self, seed: int = 42) -> None:
+        self._seed = seed
+
+    def randn(
+        self, shape: tuple[int, ...], device: str, dtype: torch.dtype
+    ) -> torch.Tensor:
+        torch.manual_seed(self._seed)
+        tensor = torch.randn(shape, device="cpu", dtype=dtype).to(device)
+        seed_and_broadcast(self._seed, tensor)
+        self._seed += 1
+        return tensor
+
+
+def _resolve_nested_attr(obj: Any, path: str) -> Any:
+    current = obj
+    for name in path.split("."):
+        if current is None or not hasattr(current, name):
+            return None
+        current = getattr(current, name)
+    return current
+
+
+def _read_config_value(model: nn.Module, keys: list[str], default: int) -> int:
+    config = getattr(model, "config", None)
+    for key in keys:
+        for root in (model, config):
+            value = _resolve_nested_attr(root, key) if root is not None else None
+            if isinstance(value, int) and value > 0:
+                return value
+    return default
+
+
+def _forward_parameter_names(module: nn.Module) -> set[str]:
+    return set(inspect.signature(module.forward).parameters.keys())
+
+
+def _infer_transformer_layout(param_names: set[str]) -> str:
+    if "img_shapes" in param_names or "txt_seq_lens" in param_names:
+        return "token_shapes"
+    if "img_ids" in param_names or "txt_ids" in param_names:
+        return "token_ids"
+    if "x" in param_names or "cap_feats" in param_names:
+        return "alias"
+    return "video"
+
+
+def _build_position_ids(
+    height: int, width: int, dims: int, device: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    img_len = height * width
+    txt_len = DEFAULT_TEXT_SEQ_LEN
+    if dims == 4:
+        img_ids = torch.zeros(img_len, 4, device=device, dtype=torch.bfloat16)
+        img_ids[:, 1] = torch.arange(height).repeat_interleave(width)
+        img_ids[:, 2] = torch.arange(width).repeat(height)
+        txt_ids = torch.zeros(txt_len, 4, device=device, dtype=torch.bfloat16)
+    else:
+        img_ids = torch.zeros(img_len, 3, device=device, dtype=torch.bfloat16)
+        img_ids[:, 0] = torch.arange(height).repeat_interleave(width)
+        img_ids[:, 1] = torch.arange(width).repeat(height)
+        txt_ids = torch.zeros(txt_len, 3, device=device, dtype=torch.bfloat16)
+    return img_ids, txt_ids
+
+
+def _build_alias_rotary_freqs(
+    model: nn.Module, device: str, height: int, width: int
+) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+    cap_len = DEFAULT_TEXT_SEQ_LEN
+    cap_pad_len = (-cap_len) % ALIAS_ROTARY_TEXT_PAD_MULTIPLE
+    cap_ids = (
+        torch.stack(
+            torch.meshgrid(
+                torch.arange(cap_len + cap_pad_len),
+                torch.arange(1),
+                torch.arange(1),
+                indexing="ij",
+            ),
+            dim=-1,
+        )
+        .flatten(0, 2)
+        .to(device)
+    )
+    img_ids = (
+        torch.stack(
+            torch.meshgrid(
+                torch.arange(1),
+                torch.arange(height // 2),
+                torch.arange(width // 2),
+                indexing="ij",
+            ),
+            dim=-1,
+        )
+        .flatten(0, 2)
+        .to(device)
+    )
+    cos_cap, sin_cap = model.rotary_emb(cap_ids)
+    cos_img, sin_img = model.rotary_emb(img_ids)
+    return ((cos_cap, sin_cap), (cos_img, sin_img))
+
+
+def _supports_image_conditioning(module: nn.Module) -> bool:
+    image_embedder = _resolve_nested_attr(module, "condition_embedder.image_embedder")
+    if image_embedder is not None:
+        return True
+    image_dim = _read_config_value(
+        module, ["arch_config.image_dim", "image_dim"], default=0
+    )
+    return image_dim > 0
+
+
+def _build_transformer_hook_inputs(
+    case: Any, model: nn.Module, device: str, ref_model: Optional[nn.Module] = None
+) -> Inputs:
+    """Build one synthetic input bundle that both transformer variants can consume."""
+    compat = _resolve_transformer_hook_compat(case)
+    param_names = _forward_parameter_names(model)
+    if ref_model is not None:
+        # The input bundle has to satisfy both call signatures.
+        param_names.update(_forward_parameter_names(ref_model))
+
+    rng = _DeterministicRNG()
+    layout = _infer_transformer_layout(param_names)
+    in_channels = _read_config_value(
+        model,
+        [
+            "arch_config.in_channels",
+            "in_channels",
+            "transformer_config.in_channels",
+        ],
+        default=DEFAULT_TRANSFORMER_IN_CHANNELS,
+    )
+    text_channels = _read_config_value(
+        model,
+        [
+            "text_states_dim",
+            "arch_config.cap_feat_dim",
+            "cap_feat_dim",
+            "caption_channels",
+            "arch_config.text_dim",
+            "text_dim",
+            "arch_config.text_embed_dim",
+            "text_embed_dim",
+            "arch_config.joint_attention_dim",
+            "joint_attention_dim",
+            "cross_attention_dim",
+            "hidden_size",
+            "dim",
+        ],
+        default=DEFAULT_TRANSFORMER_TEXT_CHANNELS,
+    )
+    pooled_channels = _read_config_value(
+        model,
+        [
+            "text_states_dim_2",
+            "arch_config.pooled_projection_dim",
+            "pooled_projection_dim",
+            "pooled_embed_dim",
+            "text_embed_dim",
+            "projection_dim",
+        ],
+        default=DEFAULT_TRANSFORMER_POOLED_CHANNELS,
+    )
+    image_channels = _read_config_value(
+        model,
+        ["arch_config.image_dim", "image_dim", "cross_attention_dim"],
+        default=I2V_IMAGE_DIM,
+    )
+
+    if layout == "token_shapes":
+        height, width = DEFAULT_TOKEN_LAYOUT_SIZE, DEFAULT_TOKEN_LAYOUT_SIZE
+        seq_len = (height // 2) * (width // 2)
+        hidden_states = rng.randn((1, seq_len, in_channels), device, torch.bfloat16)
+    elif layout == "token_ids":
+        height, width = REDUCED_TOKEN_LAYOUT_SIZE, REDUCED_TOKEN_LAYOUT_SIZE
+        seq_len = height * width
+        hidden_states = rng.randn((1, seq_len, in_channels), device, torch.bfloat16)
+    elif layout == "alias":
+        height, width = DEFAULT_TOKEN_LAYOUT_SIZE, DEFAULT_TOKEN_LAYOUT_SIZE
+        hidden_states = rng.randn(
+            (1, in_channels, 1, height, width), device, torch.bfloat16
+        )
+    elif compat.use_2d_hidden_states:
+        spatial_size = (
+            REDUCED_TOKEN_LAYOUT_SIZE
+            if "encoder_attention_mask" in param_names
+            or "encoder_hidden_states_mask" in param_names
+            else DEFAULT_TOKEN_LAYOUT_SIZE
+        )
+        height, width = spatial_size, spatial_size
+        hidden_states = rng.randn(
+            (1, in_channels, height, width),
+            device,
+            torch.bfloat16,
+        )
+    else:
+        spatial_size = (
+            REDUCED_TOKEN_LAYOUT_SIZE
+            if "encoder_attention_mask" in param_names
+            or "encoder_hidden_states_mask" in param_names
+            else DEFAULT_TOKEN_LAYOUT_SIZE
+        )
+        height, width = spatial_size, spatial_size
+        hidden_states = rng.randn(
+            (1, in_channels, DEFAULT_VIDEO_FRAME_COUNT, height, width),
+            device,
+            torch.bfloat16,
+        )
+
+    inputs: Inputs = {
+        "hidden_states": hidden_states,
+        "encoder_hidden_states": rng.randn(
+            (1, DEFAULT_TEXT_SEQ_LEN, text_channels), device, torch.bfloat16
+        ),
+        "timestep": torch.tensor(
+            [DEFAULT_TIMESTEP], device=device, dtype=torch.bfloat16
+        ),
+        "guidance": torch.tensor([1.0], device=device, dtype=torch.bfloat16),
+    }
+
+    if "pooled_projections" in param_names:
+        inputs["pooled_projections"] = rng.randn(
+            (1, pooled_channels), device, torch.bfloat16
+        )
+    if (
+        "encoder_attention_mask" in param_names
+        or "encoder_hidden_states_mask" in param_names
+    ):
+        attention_mask = torch.ones(
+            1, DEFAULT_TEXT_SEQ_LEN, device=device, dtype=torch.bool
+        )
+        inputs["encoder_attention_mask"] = attention_mask
+        inputs["encoder_hidden_states_mask"] = attention_mask
+    if "encoder_hidden_states_image" in param_names and _supports_image_conditioning(
+        model
+    ):
+        inputs["encoder_hidden_states_image"] = rng.randn(
+            (1, DEFAULT_IMAGE_TOKEN_COUNT, image_channels), device, torch.bfloat16
+        )
+    if "additional_t_cond" in param_names:
+        inputs["additional_t_cond"] = torch.zeros((1,), device=device, dtype=torch.long)
+    if "img_shapes" in param_names:
+        inputs["img_shapes"] = [[(1, height // 2, width // 2)]]
+    if "txt_seq_lens" in param_names:
+        inputs["txt_seq_lens"] = [DEFAULT_TEXT_SEQ_LEN]
+    if "img_ids" in param_names or "txt_ids" in param_names:
+        id_dims = 4 if in_channels >= LARGE_CHANNEL_LAYOUT_THRESHOLD else 3
+        img_ids, txt_ids = _build_position_ids(height, width, id_dims, device)
+        inputs["img_ids"] = img_ids
+        inputs["txt_ids"] = txt_ids
+
+    if "freqs_cis" in param_names and hasattr(model, "rotary_emb"):
+        if "img_shapes" in inputs and "txt_seq_lens" in inputs:
+            img_freqs, txt_freqs = model.rotary_emb(
+                inputs["img_shapes"],
+                inputs["txt_seq_lens"],
+                device=hidden_states.device,
+            )
+            if torch.is_complex(img_freqs) and torch.is_complex(txt_freqs):
+                inputs["freqs_cis"] = (
+                    torch.cat([img_freqs.real.float(), img_freqs.imag.float()], dim=-1),
+                    torch.cat([txt_freqs.real.float(), txt_freqs.imag.float()], dim=-1),
+                )
+            else:
+                inputs["freqs_cis"] = (img_freqs, txt_freqs)
+        elif "img_ids" in inputs and "txt_ids" in inputs:
+            ids = torch.cat([inputs["txt_ids"], inputs["img_ids"]], dim=0)
+            inputs["freqs_cis"] = model.rotary_emb(ids)
+        elif inputs["hidden_states"].ndim == 5:
+            inputs["freqs_cis"] = _build_alias_rotary_freqs(
+                model, device, height, width
+            )
+
+    inputs["hook_compat"] = compat
+    return inputs
+
+
+def _get_transformer_hook_compat(inputs: Inputs) -> TransformerHookCompat:
+    compat = inputs.get("hook_compat")
+    assert isinstance(compat, TransformerHookCompat)
+    return compat
+
+
+def _supports_guidance_embedding(module: nn.Module) -> bool:
+    time_text_embed = getattr(module, "time_text_embed", None)
+    if time_text_embed is None:
+        return True
+
+    parameters = list(inspect.signature(time_text_embed.forward).parameters.values())
+
+    if any(param.kind is inspect.Parameter.VAR_POSITIONAL for param in parameters):
+        return True
+
+    accepted_args = [
+        param
+        for param in parameters
+        if param.name != "self"
+        and param.kind
+        in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        )
+    ]
+    return len(accepted_args) >= 3
+
+
+def _prepare_transformer_hook_call(
+    module: nn.Module, inputs: Inputs, side: str
+) -> HookCall:
+    param_names = _forward_parameter_names(module)
+    signature = inspect.signature(module.forward)
+    compat = _get_transformer_hook_compat(inputs)
+    kwargs: Dict[str, Any] = {}
+    negate_output = side == "reference" and compat.negate_reference_output
+
+    if "hidden_states" in param_names:
+        kwargs["hidden_states"] = inputs["hidden_states"]
+    if "x" in param_names:
+        kwargs["x"] = [inputs["hidden_states"].squeeze(0)]
+    if "encoder_hidden_states" in param_names:
+        encoder_value: Any = inputs["encoder_hidden_states"]
+        if (
+            side == "sglang"
+            and "pooled_projections" in inputs
+            and "pooled_projections" not in param_names
+            and "encoder_attention_mask" not in param_names
+        ):
+            encoder_value = [
+                inputs["encoder_hidden_states"],
+                inputs["pooled_projections"],
+            ]
+        kwargs["encoder_hidden_states"] = encoder_value
+    if "cap_feats" in param_names:
+        kwargs["cap_feats"] = [inputs["encoder_hidden_states"].squeeze(0)]
+
+    if "timestep" in param_names:
+        timestep = inputs["timestep"]
+        if side == "reference" and compat.normalize_reference_timestep:
+            timestep = timestep / TIMESTEP_NORMALIZATION_FACTOR
+        kwargs["timestep"] = timestep
+    if "t" in param_names:
+        timestep = inputs["timestep"]
+        if side == "reference" and compat.normalize_reference_timestep:
+            timestep = timestep / TIMESTEP_NORMALIZATION_FACTOR
+        kwargs["t"] = timestep
+
+    if "guidance" in param_names and "guidance" in inputs:
+        if side == "reference" and compat.omit_reference_guidance:
+            pass
+        else:
+            skip_guidance_for_image_context = (
+                "encoder_hidden_states_image" in param_names
+                and "img_ids" not in param_names
+                and "img_shapes" not in param_names
+            )
+            supports_guidance_embedding = _supports_guidance_embedding(module)
+            requires_guidance_arg = (
+                signature.parameters["guidance"].default is inspect._empty
+            )
+            should_include_guidance = (
+                not skip_guidance_for_image_context and supports_guidance_embedding
+            )
+            if should_include_guidance or requires_guidance_arg:
+                guidance_value = inputs["guidance"]
+                if side == "sglang":
+                    guidance_value = guidance_value * TIMESTEP_NORMALIZATION_FACTOR
+                kwargs["guidance"] = guidance_value
+
+    if (
+        "encoder_hidden_states_image" in param_names
+        and "encoder_hidden_states_image" in inputs
+    ):
+        value = inputs["encoder_hidden_states_image"]
+        kwargs["encoder_hidden_states_image"] = [value] if side == "sglang" else value
+
+    for key in (
+        "pooled_projections",
+        "img_ids",
+        "txt_ids",
+        "img_shapes",
+        "txt_seq_lens",
+        "freqs_cis",
+        "additional_t_cond",
+        "encoder_attention_mask",
+        "encoder_hidden_states_mask",
+    ):
+        if key in param_names and key in inputs:
+            kwargs[key] = inputs[key]
+
+    if "return_dict" in param_names:
+        kwargs["return_dict"] = True
+
+    return HookCall(module=module, kwargs=kwargs, negate_output=negate_output)
+
+
+def _prepare_transformer_sglang_call(module: nn.Module, inputs: Inputs) -> HookCall:
+    return _prepare_transformer_hook_call(module, inputs, side="sglang")
+
+
+def _prepare_transformer_reference_call(module: nn.Module, inputs: Inputs) -> HookCall:
+    return _prepare_transformer_hook_call(module, inputs, side="reference")
+
+
+class _VAEDecodeModule(nn.Module):
+    def __init__(self, vae: nn.Module):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        if (
+            any(
+                isinstance(module, (nn.Conv3d, nn.ConvTranspose3d))
+                for module in self.vae.modules()
+            )
+            and z.ndim == 4
+        ):
+            z = z.unsqueeze(2)
+        output = self.vae.decode(z)
+        tensor = output.sample if hasattr(output, "sample") else output
+        if isinstance(tensor, (list, tuple)):
+            tensor = tensor[0]
+        return tensor.squeeze(2) if tensor.ndim == 5 else tensor
+
+
+def _infer_vae_latent_channels(model: nn.Module) -> int:
+    for path in ("post_quant_conv.in_channels", "post_quant_conv.conv.in_channels"):
+        value = _resolve_nested_attr(model, path)
+        if isinstance(value, int) and value > 0:
+            return value
+    return _read_config_value(
+        model,
+        [
+            "z_dim",
+            "arch_config.z_dim",
+            "latent_channels",
+            "arch_config.latent_channels",
+            "num_channels_latents",
+            "arch_config.num_channels_latents",
+            "latent_dim",
+            "z_channels",
+            "arch_config.z_channels",
+        ],
+        default=DEFAULT_VAE_LATENT_CHANNELS,
+    )
+
+
+def _build_vae_hook_inputs(
+    case: Any, model: nn.Module, device: str, ref_model: Optional[nn.Module] = None
+) -> Inputs:
+    del case, ref_model
+    latent_channels = _infer_vae_latent_channels(model)
+    rng = _DeterministicRNG()
+    return {
+        "z": rng.randn(
+            (
+                1,
+                latent_channels,
+                DEFAULT_VAE_LATENT_SPATIAL_SIZE,
+                DEFAULT_VAE_LATENT_SPATIAL_SIZE,
+            ),
+            device,
+            torch.bfloat16,
+        )
+    }
+
+
+def _prepare_vae_decode_call(module: nn.Module, inputs: Inputs) -> HookCall:
+    return HookCall(module=_VAEDecodeModule(module), args=(inputs["z"],))
+
+
+TRANSFORMER_NATIVE_PROFILE = NativeHookProfile(
+    build_inputs=_build_transformer_hook_inputs,
+    prepare_sglang_call=_prepare_transformer_sglang_call,
+    prepare_reference_call=_prepare_transformer_reference_call,
+)
+
+VAE_NATIVE_PROFILE = NativeHookProfile(
+    build_inputs=_build_vae_hook_inputs,
+    prepare_sglang_call=_prepare_vae_decode_call,
+    prepare_reference_call=_prepare_vae_decode_call,
+)
+
+
+def resolve_component_native_profile(component: ComponentType) -> NativeHookProfile:
+    if component == ComponentType.TRANSFORMER:
+        return TRANSFORMER_NATIVE_PROFILE
+    if component == ComponentType.VAE:
+        return VAE_NATIVE_PROFILE
+    raise KeyError(f"Unsupported native accuracy component: {component.value}")

python/sglang/multimodal_gen/test/server/accuracy_utils.py

@@ -0,0 +1,880 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file as safetensors_load_file
+from torch.distributed.tensor import distribute_tensor
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    destroy_model_parallel,
+    get_data_parallel_world_size,
+    get_sequence_parallel_world_size,
+    get_tensor_model_parallel_world_size,
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    DEFAULT_TEXT_ENCODER_VOCAB_SIZE,
+    I2V_TEXT_ENCODER_DIM,
+    TEXT_ENCODER_INPUT_SEED,
+    TEXT_ENCODER_TOKEN_LENGTH,
+    TEXT_ENCODER_TOKEN_MAX,
+    TEXT_ENCODER_TOKEN_MIN,
+    ComponentType,
+    get_threshold,
+)
+
+STAGED_1GPU_NATIVE_CASE_IDS = {
+    "flux_2_image_t2i",
+    "qwen_image_layered_i2i",
+    "flux_2_image_t2i_upscaling_4x",
+    "flux_2_ti2i",
+    "flux_2_t2i_customized_vae_path",
+    "flux_2_ti2i_multi_image_cache_dit",
+}
+
+# These case allowlists are accuracy-runner policy. They select the few 1-GPU
+# cases that need sequential SGLang/reference execution to stay within memory
+# limits during CI and local correctness runs.
+STAGED_1GPU_TEXT_ENCODER_CASE_IDS = {
+    "flux_2_image_t2i",
+    "flux_2_image_t2i_upscaling_4x",
+    "mova_360p_1gpu",
+    "flux_2_ti2i",
+    "flux_2_t2i_customized_vae_path",
+    "flux_2_ti2i_multi_image_cache_dit",
+}
+
+SOURCE_PREFIXES = (
+    "module.",
+    "model.",
+    "transformer.",
+    "text_encoder.",
+    "image_encoder.",
+    "encoder.",
+    "decoder.",
+    "model.language_model.",
+    "model.visual.",
+)
+
+TARGET_PREFIXES = (
+    "module.",
+    "model.",
+    "transformer.",
+    "text_encoder.",
+    "image_encoder.",
+    "encoder.",
+    "decoder.",
+)
+
+
+@dataclass(frozen=True)
+class ComponentSelection:
+    base_model_id: str
+    base_model_root: str
+    component_paths: Dict[str, str]
+    source_root: str
+    source_path: str
+    source_subfolder: str
+
+
+@dataclass(frozen=True)
+class ParameterShardContext:
+    world_size: int
+    rank: int
+
+
+def seed_and_broadcast(seed: int, tensor: torch.Tensor) -> torch.Tensor:
+    """Seed and broadcast tensor across ranks for determinism."""
+    torch.manual_seed(seed)
+    if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
+        torch.distributed.broadcast(tensor, src=0)
+    return tensor
+
+
+def read_json_file(path: str) -> Dict[str, Any]:
+    if not os.path.exists(path):
+        return {}
+    with open(path) as f:
+        return json.load(f)
+
+
+def has_component_files(path: str) -> bool:
+    if not os.path.isdir(path):
+        return False
+    if os.path.exists(os.path.join(path, "config.json")):
+        return True
+    for ext in (".safetensors", ".bin", ".pth"):
+        if any(name.endswith(ext) for name in os.listdir(path)):
+            return True
+    return False
+
+
+def list_safetensor_files(path: str) -> List[str]:
+    if not os.path.isdir(path):
+        return []
+    return sorted(
+        os.path.join(path, name)
+        for name in os.listdir(path)
+        if name.endswith(".safetensors")
+    )
+
+
+def is_text_encoder_config(path: str) -> bool:
+    cfg_path = os.path.join(path, "config.json")
+    if not os.path.exists(cfg_path):
+        return False
+    cfg = read_json_file(cfg_path)
+    if cfg.get("model_type") == "i2v" or cfg.get("dim") == I2V_TEXT_ENCODER_DIM:
+        return False
+    return True
+
+
+def _resolve_component_subfolder(
+    model_index: Dict[str, Any], key: str
+) -> Optional[str]:
+    entry = model_index.get(key)
+    if isinstance(entry, dict):
+        return entry.get("path") or entry.get("subfolder")
+    if isinstance(entry, str):
+        return entry
+    if entry is not None:
+        return key
+    return None
+
+
+def resolve_component_path(
+    local_root: str, component: ComponentType, model_index_keys: Tuple[str, ...]
+) -> Tuple[str, str]:
+    model_index_path = os.path.join(local_root, "model_index.json")
+    model_index = read_json_file(model_index_path)
+
+    if model_index:
+        for key in model_index_keys:
+            subfolder = _resolve_component_subfolder(model_index, key)
+            if not subfolder:
+                continue
+            candidate = os.path.join(local_root, subfolder)
+            if not has_component_files(candidate):
+                continue
+            if component == ComponentType.TEXT_ENCODER and not is_text_encoder_config(
+                candidate
+            ):
+                continue
+            return candidate, subfolder
+
+    if has_component_files(local_root):
+        if component != ComponentType.TEXT_ENCODER or is_text_encoder_config(
+            local_root
+        ):
+            return local_root, ""
+
+    raise FileNotFoundError(
+        f"Could not resolve {component.value} from model_index.json under {local_root}"
+    )
+
+
+def extract_component_path_overrides(extra_args: List[str]) -> Dict[str, str]:
+    component_paths: Dict[str, str] = {}
+    index = 0
+    while index < len(extra_args):
+        arg = extra_args[index]
+        key_part = arg.split("=", 1)[0] if "=" in arg else arg
+        if key_part.startswith("--") and key_part.endswith("-path"):
+            component = key_part[2:-5].replace("-", "_")
+            if "=" in arg:
+                component_paths[component] = arg.split("=", 1)[1]
+            elif index + 1 < len(extra_args) and not extra_args[index + 1].startswith(
+                "-"
+            ):
+                index += 1
+                component_paths[component] = extra_args[index]
+        index += 1
+
+    for component, path in component_paths.items():
+        component_paths[component] = os.path.expanduser(path)
+    return component_paths
+
+
+def load_checkpoint_weights(
+    module: nn.Module, model_path: str
+) -> tuple[list[str], list[str]]:
+    safetensors_files = list_safetensor_files(model_path)
+    assert safetensors_files, f"Found no safetensors files in {model_path}"
+
+    loaded_state: Dict[str, torch.Tensor] = {}
+    for safetensor_path in safetensors_files:
+        loaded_state.update(safetensors_load_file(safetensor_path))
+
+    module.load_state_dict(loaded_state, strict=False)
+
+    state_keys = set(module.state_dict().keys())
+    loaded_keys = set(loaded_state.keys())
+    missing_keys = sorted(state_keys - loaded_keys)
+    unexpected_keys = sorted(loaded_keys - state_keys)
+    return missing_keys, unexpected_keys
+
+
+def select_component_source(
+    model_id: str,
+    extra_args: List[str],
+    component: ComponentType,
+    model_index_keys: Tuple[str, ...],
+) -> ComponentSelection:
+    component_paths = extract_component_path_overrides(extra_args)
+    base_model_root = maybe_download_model(model_id)
+    search_keys = [component.value]
+    for key in model_index_keys:
+        if key not in search_keys:
+            search_keys.append(key)
+
+    source_root = base_model_root
+    component_key = component.value
+    for key in search_keys:
+        override_path = component_paths.get(key)
+        if override_path:
+            source_root = maybe_download_model(override_path)
+            component_key = key
+            break
+
+    ordered_keys = [component_key]
+    for key in search_keys:
+        if key not in ordered_keys:
+            ordered_keys.append(key)
+    source_path, source_subfolder = resolve_component_path(
+        source_root,
+        component,
+        tuple(ordered_keys),
+    )
+    return ComponentSelection(
+        base_model_id=model_id,
+        base_model_root=base_model_root,
+        component_paths=component_paths,
+        source_root=source_root,
+        source_path=source_path,
+        source_subfolder=source_subfolder,
+    )
+
+
+def ensure_distributed_env_defaults() -> None:
+    if "WORLD_SIZE" in os.environ:
+        return
+    os.environ.update(
+        {
+            "MASTER_ADDR": os.getenv("MASTER_ADDR", "127.0.0.1"),
+            "MASTER_PORT": os.getenv("MASTER_PORT", "29505"),
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+        }
+    )
+
+
+def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
+    tp_size = sgl_args.tp_size
+    sp_degree = sgl_args.sp_degree
+    ulysses_degree = sgl_args.ulysses_degree
+    ring_degree = sgl_args.ring_degree
+    dp_size = sgl_args.dp_size
+    enable_cfg_parallel = bool(sgl_args.enable_cfg_parallel)
+
+    if (
+        tp_size is None
+        or sp_degree is None
+        or ulysses_degree is None
+        or ring_degree is None
+    ):
+        raise RuntimeError(
+            "ServerArgs must have tp_size, sp_degree, ulysses_degree, and ring_degree before init"
+        )
+
+    if not model_parallel_is_initialized() and torch.distributed.is_initialized():
+        # A prior case may have failed while distributed groups were only partially
+        # initialized. Clear any stale group objects before re-initializing.
+        destroy_model_parallel()
+
+    if model_parallel_is_initialized():
+        current_tp = get_tensor_model_parallel_world_size()
+        current_sp = get_sequence_parallel_world_size()
+        current_dp = get_data_parallel_world_size()
+        if current_tp == tp_size and current_sp == sp_degree and current_dp == dp_size:
+            return
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+        destroy_model_parallel()
+
+    ensure_distributed_env_defaults()
+
+    maybe_init_distributed_environment_and_model_parallel(
+        tp_size=tp_size,
+        sp_size=sp_degree,
+        enable_cfg_parallel=enable_cfg_parallel,
+        ulysses_degree=ulysses_degree,
+        ring_degree=ring_degree,
+        dp_size=dp_size,
+    )
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+
+
+def build_accuracy_server_args(
+    base_model_id: str,
+    base_model_root: str,
+    case: Any,
+    component: ComponentType,
+    num_gpus: int,
+    component_paths: Dict[str, str],
+) -> ServerArgs:
+    cfg_parallel = bool(case.server_args.cfg_parallel)
+    kwargs = {
+        "model_path": base_model_root,
+        "model_id": base_model_id,
+        "num_gpus": num_gpus,
+        "trust_remote_code": True,
+        "component_paths": component_paths,
+        "enable_cfg_parallel": cfg_parallel,
+    }
+
+    if case.server_args.tp_size is not None:
+        kwargs["tp_size"] = case.server_args.tp_size
+    if case.server_args.ulysses_degree is not None:
+        kwargs["ulysses_degree"] = case.server_args.ulysses_degree
+    if case.server_args.ring_degree is not None:
+        kwargs["ring_degree"] = case.server_args.ring_degree
+
+    if component == ComponentType.TEXT_ENCODER:
+        kwargs["enable_cfg_parallel"] = False
+
+    sgl_args = ServerArgs.from_kwargs(**kwargs)
+    sgl_args.text_encoder_cpu_offload = False
+    sgl_args.dit_cpu_offload = False
+    sgl_args.vae_cpu_offload = False
+    sgl_args.image_encoder_cpu_offload = False
+    sgl_args.enable_cache_dit = case.server_args.enable_cache_dit
+    sgl_args.dit_layerwise_offload = case.server_args.dit_layerwise_offload
+    sgl_args.dit_offload_prefetch_size = case.server_args.dit_offload_prefetch_size
+    return sgl_args
+
+
+def set_module_attr(module: nn.Module, name: str, value: Any) -> None:
+    """Assign to a nested parameter/buffer path such as `blocks.0.attn.to_q.weight`."""
+    attrs = name.split(".")
+    parent = module
+    for attr in attrs[:-1]:
+        if hasattr(parent, attr):
+            parent = getattr(parent, attr)
+        elif isinstance(parent, (nn.ModuleList, nn.Sequential)):
+            parent = parent[int(attr)]
+        elif isinstance(parent, nn.ModuleDict):
+            parent = parent[attr]
+        else:
+            raise AttributeError(
+                f"Cannot resolve {name} on {module.__class__.__name__}"
+            )
+    setattr(parent, attrs[-1], value)
+
+
+def materialize_module(
+    module: nn.Module, device: torch.device, dtype: torch.dtype
+) -> None:
+    """Materialize meta tensors and cast floating tensors onto one target device/dtype."""
+    for name, param in module.named_parameters():
+        if param.device.type == "meta":
+            new_data = torch.zeros(param.shape, device=device, dtype=dtype)
+            if hasattr(param, "device_mesh") and param.device_mesh is not None:
+                new_data = distribute_tensor(
+                    new_data, param.device_mesh, param.placements
+                )
+            set_module_attr(
+                module, name, nn.Parameter(new_data, requires_grad=param.requires_grad)
+            )
+        elif torch.is_floating_point(param):
+            param.data = param.data.to(device=device, dtype=dtype)
+
+    for name, buf in module.named_buffers():
+        if buf.device.type == "meta":
+            new_buf = torch.zeros(buf.shape, device=device, dtype=buf.dtype)
+            if hasattr(buf, "device_mesh") and buf.device_mesh is not None:
+                new_buf = distribute_tensor(new_buf, buf.device_mesh, buf.placements)
+            set_module_attr(module, name, new_buf)
+        elif torch.is_floating_point(buf):
+            buf.data = buf.data.to(device=device, dtype=dtype)
+
+
+def build_parameter_shard_contexts(
+    module: nn.Module,
+) -> Dict[str, ParameterShardContext]:
+    """Record TP shard world/rank for each parameter owned by a TP-aware submodule."""
+    shard_contexts: Dict[str, ParameterShardContext] = {}
+    for module_name, submodule in module.named_modules():
+        tp_group = getattr(submodule, "tp_group", None)
+        if tp_group is None:
+            continue
+
+        context = ParameterShardContext(
+            world_size=get_group_size(tp_group),
+            rank=get_group_rank(tp_group),
+        )
+        if context.world_size <= 1:
+            continue
+
+        for name, _ in submodule.named_parameters(recurse=False):
+            qualified_name = f"{module_name}.{name}" if module_name else name
+            shard_contexts[qualified_name] = context
+        for name, _ in submodule.named_buffers(recurse=False):
+            qualified_name = f"{module_name}.{name}" if module_name else name
+            shard_contexts[qualified_name] = context
+
+    return shard_contexts
+
+
+def build_state_lookup(state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Index a source state dict under both original and prefix-stripped names."""
+    lookup: Dict[str, torch.Tensor] = {}
+    for key, val in state.items():
+        lookup[key] = val
+        for prefix in SOURCE_PREFIXES:
+            if key.startswith(prefix):
+                lookup[key[len(prefix) :]] = val
+    return lookup
+
+
+def normalize_state_key(name: str) -> str:
+    """Normalize common naming differences between source and target state dicts."""
+    return (
+        name.replace("_fsdp_wrapped_module.", "")
+        .replace("_orig_mod.", "")
+        .replace("gamma", "weight")
+        .replace("beta", "bias")
+        .replace("scale", "weight")
+        .replace("shift", "bias")
+    )
+
+
+def fuse_qkv(lookup: Dict[str, torch.Tensor], name: str) -> Optional[torch.Tensor]:
+    if "qkv_proj" not in name:
+        return None
+    variants = ["q_proj", "q"]
+    for repl in variants:
+        q_name = name.replace("qkv_proj", repl)
+        k_name = q_name.replace(".q_proj", ".k_proj").replace(".q", ".k")
+        v_name = q_name.replace(".q_proj", ".v_proj").replace(".q", ".v")
+        if q_name in lookup and k_name in lookup and v_name in lookup:
+            return torch.cat([lookup[q_name], lookup[k_name], lookup[v_name]], dim=0)
+    return None
+
+
+def fuse_gate_up_proj(
+    lookup: Dict[str, torch.Tensor], name: str
+) -> Optional[torch.Tensor]:
+    if "gate_up_proj" not in name:
+        return None
+
+    for gate_token, up_token in (("gate_proj", "up_proj"), ("wi_0", "wi_1")):
+        gate_name = name.replace("gate_up_proj", gate_token)
+        up_name = name.replace("gate_up_proj", up_token)
+        if gate_name in lookup and up_name in lookup:
+            return torch.cat([lookup[gate_name], lookup[up_name]], dim=0)
+    return None
+
+
+def generate_name_candidates(
+    name: str, reverse_mapping: Optional[Dict[str, Tuple[str, Any, Any]]]
+) -> List[str]:
+    candidates: List[str] = []
+    clean = normalize_state_key(name)
+
+    for cand in (name, clean):
+        if cand not in candidates:
+            candidates.append(cand)
+
+    if reverse_mapping:
+        for key in (name, clean):
+            entry = reverse_mapping.get(key)
+            if entry and entry[0] not in candidates:
+                candidates.append(entry[0])
+
+    for prefix in TARGET_PREFIXES:
+        if clean.startswith(prefix):
+            stripped = clean[len(prefix) :]
+            if stripped and stripped not in candidates:
+                candidates.append(stripped)
+
+    parts = clean.split(".")
+    for i in range(1, len(parts)):
+        cand = ".".join(parts[i:])
+        if cand not in candidates:
+            candidates.append(cand)
+
+    return candidates
+
+
+def copy_tensor(
+    dest: torch.Tensor,
+    src: torch.Tensor,
+    tp_world: int,
+    rank: int,
+) -> bool:
+    if src.numel() == 0:
+        return False
+    src = src.to(device=dest.device, dtype=dest.dtype)
+
+    if hasattr(dest, "device_mesh") and dest.device_mesh is not None:
+        if src.numel() == dest.numel():
+            with torch.no_grad():
+                dt = distribute_tensor(
+                    src.view(dest.shape), dest.device_mesh, dest.placements
+                )
+                dest.copy_(dt)
+            return True
+
+    if src.numel() == dest.numel():
+        with torch.no_grad():
+            dest.copy_(src.view(dest.shape))
+        return True
+
+    if tp_world > 1 and src.numel() == dest.numel() * tp_world:
+        if (
+            src.ndim == dest.ndim
+            and src.shape[0] == dest.shape[0] * tp_world
+            and src.shape[1:] == dest.shape[1:]
+        ):
+            with torch.no_grad():
+                dest.copy_(src[rank * dest.shape[0] : (rank + 1) * dest.shape[0], ...])
+            return True
+        if (
+            src.ndim >= 2
+            and dest.ndim >= 2
+            and src.ndim == dest.ndim
+            and src.shape[0] == dest.shape[0]
+            and src.shape[1] == dest.shape[1] * tp_world
+            and src.shape[2:] == dest.shape[2:]
+        ):
+            with torch.no_grad():
+                dest.copy_(src[:, rank * dest.shape[1] : (rank + 1) * dest.shape[1]])
+            return True
+
+    if src.ndim == 4 and dest.ndim == 5 and dest.numel() == src.numel() * dest.shape[2]:
+        with torch.no_grad():
+            dest.copy_(src.unsqueeze(2).repeat(1, 1, dest.shape[2], 1, 1))
+        return True
+
+    return False
+
+
+def _config_to_dict(config: Any) -> Dict[str, Any]:
+    to_dict = getattr(config, "to_dict", None)
+    if not callable(to_dict):
+        return {}
+    config_dict = to_dict()
+    return config_dict if isinstance(config_dict, dict) else {}
+
+
+def resolve_text_encoder_vocab_size(config: Any) -> int:
+    config_dict = _config_to_dict(config)
+    vocab_size = config_dict.get("vocab_size")
+    if isinstance(vocab_size, int) and vocab_size > 0:
+        return vocab_size
+
+    text_config = config_dict.get("text_config")
+    if isinstance(text_config, dict):
+        nested_vocab_size = text_config.get("vocab_size")
+        if isinstance(nested_vocab_size, int) and nested_vocab_size > 0:
+            return nested_vocab_size
+
+    return DEFAULT_TEXT_ENCODER_VOCAB_SIZE
+
+
+def build_deterministic_text_encoder_inputs(
+    config: Any, device: str
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build one stable token batch that works across text-encoder implementations."""
+    vocab_size = resolve_text_encoder_vocab_size(config)
+    max_token_id = max(
+        TEXT_ENCODER_TOKEN_MIN + 1, min(vocab_size, TEXT_ENCODER_TOKEN_MAX)
+    )
+
+    torch.manual_seed(TEXT_ENCODER_INPUT_SEED)
+    input_ids = torch.randint(
+        TEXT_ENCODER_TOKEN_MIN,
+        max_token_id,
+        (1, TEXT_ENCODER_TOKEN_LENGTH),
+        device="cpu",
+        dtype=torch.long,
+    ).to(device)
+    attention_mask = torch.ones_like(input_ids)
+    return input_ids, attention_mask
+
+
+def resolve_text_encoder_forward_module(model: nn.Module) -> nn.Module:
+    get_encoder = getattr(model, "get_encoder", None)
+    return get_encoder() if callable(get_encoder) else model
+
+
+def _module_device(module: nn.Module) -> torch.device:
+    param = next(module.parameters(), None)
+    if param is not None:
+        return param.device
+
+    buf = next(module.buffers(), None)
+    if buf is not None:
+        return buf.device
+
+    return torch.device("cpu")
+
+
+def extract_output_tensor(output: Any) -> torch.Tensor:
+    """Best-effort extraction of a tensor from model outputs."""
+    if isinstance(output, torch.Tensor):
+        return output
+
+    sample = getattr(output, "sample", None)
+    if sample is not None:
+        if isinstance(sample, (list, tuple)):
+            sample = sample[0]
+        if isinstance(sample, torch.Tensor):
+            return sample
+
+    last_hidden_state = getattr(output, "last_hidden_state", None)
+    if last_hidden_state is not None:
+        return last_hidden_state
+
+    hidden_states = getattr(output, "hidden_states", None)
+    if hidden_states:
+        return hidden_states[-1]
+
+    pooler_output = getattr(output, "pooler_output", None)
+    if pooler_output is not None:
+        return pooler_output
+
+    logits = getattr(output, "logits", None)
+    if logits is not None:
+        return logits
+
+    if (
+        isinstance(output, (list, tuple))
+        and output
+        and isinstance(output[0], torch.Tensor)
+    ):
+        return output[0]
+    raise ValueError(f"Could not extract tensor from output of type {type(output)}")
+
+
+def run_text_encoder_accuracy_pair(
+    sgl: nn.Module, ref: nn.Module
+) -> tuple[torch.Tensor, torch.Tensor]:
+    input_ids, attention_mask = build_deterministic_text_encoder_inputs(
+        ref.config, "cpu"
+    )
+    return (
+        _run_single_text_encoder_forward(sgl, input_ids, attention_mask),
+        _run_single_text_encoder_forward(ref, input_ids, attention_mask),
+    )
+
+
+def _should_stage_case(case: Any, component: ComponentType, num_gpus: int) -> bool:
+    if num_gpus == 2:
+        return True
+    if num_gpus != 1:
+        return False
+    if component == ComponentType.TEXT_ENCODER:
+        return case.id in STAGED_1GPU_TEXT_ENCODER_CASE_IDS
+    return case.id in STAGED_1GPU_NATIVE_CASE_IDS
+
+
+def _run_single_text_encoder_forward(
+    model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor
+) -> torch.Tensor:
+    """Run one encoder forward and normalize its output into a tensor."""
+    with torch.no_grad():
+        forward_model = resolve_text_encoder_forward_module(model)
+        model_device = _module_device(forward_model)
+        output = forward_model(
+            input_ids.to(device=model_device),
+            attention_mask=attention_mask.to(device=model_device),
+            output_hidden_states=True,
+        )
+    return extract_output_tensor(output)
+
+
+def _run_staged_native_component_accuracy_case(
+    engine_cls: Any,
+    case: Any,
+    component: ComponentType,
+    library: str,
+    num_gpus: int,
+) -> None:
+    from sglang.multimodal_gen.test.server.accuracy_hooks import (
+        resolve_component_native_profile,
+    )
+
+    sgl = None
+    ref = None
+    try:
+        sgl, ref, device = engine_cls.load_component_pair(
+            case,
+            component,
+            library,
+            num_gpus,
+            materialize_ref_on_device=False,
+        )
+        profile = resolve_component_native_profile(component)
+        inputs = profile.build_inputs(case, sgl, device, ref)
+
+        sgl_call = profile.prepare_sglang_call(sgl, inputs)
+        with torch.no_grad():
+            sgl_raw = engine_cls._execute_with_native_hook(sgl_call)
+        sgl_out = profile.normalize_sglang_output(sgl_raw)
+        sgl_out = engine_cls._apply_output_transforms(sgl_out, sgl_call).detach().cpu()
+
+        del sgl_call
+        del sgl_raw
+        del sgl
+        sgl = None
+        engine_cls.clear_memory()
+
+        ref = ref.to(device=device, dtype=torch.bfloat16).eval()
+        ref_call = profile.prepare_reference_call(ref, inputs)
+        with torch.no_grad():
+            ref_raw = engine_cls._execute_with_native_hook(ref_call)
+        ref_out = profile.normalize_reference_output(ref_raw)
+        ref_out = engine_cls._apply_output_transforms(ref_out, ref_call).detach().cpu()
+        del ref_call
+        del ref_raw
+
+        engine_cls.check_accuracy(
+            sgl_out,
+            ref_out,
+            f"{case.id}_{component.value}",
+            get_threshold(case.id, component),
+        )
+    finally:
+        if sgl is not None:
+            del sgl
+        if ref is not None:
+            del ref
+        engine_cls.reset_parallel_runtime()
+        engine_cls.clear_memory()
+
+
+def _run_staged_text_encoder_accuracy_case(
+    engine_cls: Any, case: Any, num_gpus: int
+) -> None:
+    sgl = None
+    ref = None
+    try:
+        sgl, ref, device = engine_cls.load_component_pair(
+            case,
+            ComponentType.TEXT_ENCODER,
+            "transformers",
+            num_gpus,
+            materialize_sgl_on_device=False,
+            materialize_ref_on_device=False,
+        )
+        input_ids, attention_mask = build_deterministic_text_encoder_inputs(
+            ref.config, "cpu"
+        )
+
+        sgl = sgl.to(device=device, dtype=torch.bfloat16).eval()
+        sgl_out = (
+            _run_single_text_encoder_forward(sgl, input_ids, attention_mask)
+            .detach()
+            .cpu()
+        )
+
+        del sgl
+        sgl = None
+        engine_cls.clear_memory()
+
+        ref = ref.to(device=device, dtype=torch.bfloat16).eval()
+        ref_out = (
+            _run_single_text_encoder_forward(ref, input_ids, attention_mask)
+            .detach()
+            .cpu()
+        )
+
+        engine_cls.check_accuracy(
+            sgl_out,
+            ref_out,
+            f"{case.id}_encoder",
+            get_threshold(case.id, ComponentType.TEXT_ENCODER),
+        )
+    finally:
+        if sgl is not None:
+            del sgl
+        if ref is not None:
+            del ref
+        engine_cls.reset_parallel_runtime()
+        engine_cls.clear_memory()
+
+
+def run_native_component_accuracy_case(
+    engine_cls: Any,
+    case: Any,
+    component: ComponentType,
+    library: str,
+    num_gpus: int,
+) -> None:
+    if _should_stage_case(case, component, num_gpus):
+        _run_staged_native_component_accuracy_case(
+            engine_cls, case, component, library, num_gpus
+        )
+        return
+    engine_cls.clear_memory()
+    sgl = None
+    ref = None
+    try:
+        sgl, ref, device = engine_cls.load_component_pair(
+            case, component, library, num_gpus
+        )
+        sgl_out, ref_out = engine_cls.run_component_pair_native(
+            case, component, sgl, ref, device
+        )
+        engine_cls.check_accuracy(
+            sgl_out,
+            ref_out,
+            f"{case.id}_{component.value}",
+            get_threshold(case.id, component),
+        )
+    finally:
+        if sgl is not None:
+            del sgl
+        if ref is not None:
+            del ref
+        engine_cls.reset_parallel_runtime()
+        engine_cls.clear_memory()
+
+
+def run_text_encoder_accuracy_case(engine_cls: Any, case: Any, num_gpus: int) -> None:
+    if _should_stage_case(case, ComponentType.TEXT_ENCODER, num_gpus):
+        _run_staged_text_encoder_accuracy_case(engine_cls, case, num_gpus)
+        return
+    engine_cls.clear_memory()
+    sgl = None
+    ref = None
+    try:
+        sgl, ref, _device = engine_cls.load_component_pair(
+            case, ComponentType.TEXT_ENCODER, "transformers", num_gpus
+        )
+        sgl_out, ref_out = run_text_encoder_accuracy_pair(sgl, ref)
+        engine_cls.check_accuracy(
+            sgl_out,
+            ref_out,
+            f"{case.id}_encoder",
+            get_threshold(case.id, ComponentType.TEXT_ENCODER),
+        )
+    finally:
+        if sgl is not None:
+            del sgl
+        if ref is not None:
+            del ref
+        engine_cls.reset_parallel_runtime()
+        engine_cls.clear_memory()

python/sglang/multimodal_gen/test/server/component_accuracy.py

@@ -0,0 +1,534 @@
+from __future__ import annotations
+
+import gc
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import diffusers
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    T5EncoderModel,
+    UMT5EncoderModel,
+)
+
+try:
+    from transformers import AutoModelForImageTextToText as AutoVisionTextModel
+except ImportError:
+    try:
+        from transformers import AutoModelForVision2Seq as AutoVisionTextModel
+    except ImportError:
+        AutoVisionTextModel = None
+
+import sglang.multimodal_gen.runtime.managers.forward_context as fc_mod
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    destroy_model_parallel,
+    get_local_torch_device,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    get_param_names_mapping,
+    hf_to_custom_state_dict,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import ForwardContext
+from sglang.multimodal_gen.runtime.models.vaes.wanvae import AutoencoderKLWan
+from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    DEFAULT_TIMESTEP,
+    ComponentType,
+)
+from sglang.multimodal_gen.test.server.accuracy_hooks import (
+    resolve_component_native_profile,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    build_accuracy_server_args,
+    build_parameter_shard_contexts,
+    build_state_lookup,
+    copy_tensor,
+    fuse_gate_up_proj,
+    fuse_qkv,
+    generate_name_candidates,
+    initialize_parallel_runtime,
+    load_checkpoint_weights,
+    materialize_module,
+    read_json_file,
+    resolve_text_encoder_forward_module,
+    select_component_source,
+)
+from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+
+logger = init_logger(__name__)
+
+MIN_MATCH_RATIO = float(os.getenv("SGLANG_DIFFUSION_WEIGHT_MATCH_RATIO", "0.98"))
+
+
+@dataclass(frozen=True)
+class ComponentSpec:
+    model_index_keys: Tuple[str, ...]
+    reference_library: str
+
+
+COMPONENT_SPECS: Dict[ComponentType, ComponentSpec] = {
+    ComponentType.VAE: ComponentSpec(
+        model_index_keys=(
+            "vae",
+            "vae_model",
+            "autoencoder",
+            "autoencoder_kl",
+            "video_vae",
+            "audio_vae",
+        ),
+        reference_library="diffusers",
+    ),
+    ComponentType.TRANSFORMER: ComponentSpec(
+        model_index_keys=("transformer", "unet", "dit", "video_dit", "audio_dit"),
+        reference_library="diffusers",
+    ),
+    ComponentType.TEXT_ENCODER: ComponentSpec(
+        model_index_keys=(
+            "text_encoder",
+            "text_encoder_2",
+            "text_encoder_3",
+            "image_encoder",
+        ),
+        reference_library="transformers",
+    ),
+}
+
+
+# Component loading helpers
+def _load_sglang_component(
+    comp_path: str,
+    sgl_args: ServerArgs,
+    component: ComponentType,
+    library: str,
+    text_encoder_cpu_offload: bool | None = None,
+) -> nn.Module:
+    loader = ComponentLoader.for_component_type(component.value, library)
+    if component == ComponentType.TEXT_ENCODER:
+        component_model = loader.load_customized(
+            comp_path,
+            sgl_args,
+            component.value,
+            cpu_offload_flag=text_encoder_cpu_offload,
+        )
+    else:
+        component_model = loader.load_customized(comp_path, sgl_args, component.value)
+    if component_model is None:
+        raise RuntimeError(f"Failed to load customized {component.value}")
+    return component_model
+
+
+def _load_wan_reference_vae(comp_path: str, pipeline_config) -> nn.Module:
+    vae_config = pipeline_config.vae_config
+    vae_config.update_model_arch(
+        get_diffusers_component_config(component_path=comp_path)
+    )
+    if hasattr(vae_config, "post_init"):
+        vae_config.post_init()
+
+    vae = AutoencoderKLWan(vae_config)
+    missing_keys, unexpected_keys = load_checkpoint_weights(vae, comp_path)
+    if missing_keys:
+        logger.warning("WAN VAE missing keys: %s", missing_keys)
+    if unexpected_keys:
+        logger.warning("WAN VAE unexpected keys: %s", unexpected_keys)
+    return vae
+
+
+def _load_reference_component(
+    comp_path: str,
+    source_root: str,
+    component: ComponentType,
+    hub_id: str,
+    pipeline_config,
+    subfolder: str,
+) -> nn.Module:
+    # WAN VAE does not have a clean generic diffusers auto-load path here, and we
+    # explicitly need checkpoint-loaded weights for reference-side transfer/parity.
+    if component == ComponentType.VAE and "wan" in hub_id.lower():
+        return _load_wan_reference_vae(comp_path, pipeline_config)
+
+    if component == ComponentType.VAE:
+        cfg = read_json_file(os.path.join(comp_path, "config.json"))
+        class_name = cfg.get("_class_name") if cfg else None
+        cls = getattr(diffusers, str(class_name), None) if class_name else None
+        if cls is None:
+            cls = diffusers.AutoencoderKL
+        return cls.from_pretrained(
+            source_root,
+            subfolder=subfolder,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        )
+
+    if component == ComponentType.TRANSFORMER:
+        cfg = read_json_file(os.path.join(comp_path, "config.json"))
+        class_name = cfg.get("_class_name") if cfg else None
+        load_kwargs: Dict[str, Any] = {
+            "torch_dtype": torch.bfloat16,
+            "trust_remote_code": True,
+        }
+        if cfg:
+            for k, out_k in [
+                ("in_dim", "in_channels"),
+                ("dim", "hidden_size"),
+                ("num_heads", "num_attention_heads"),
+                ("out_dim", "out_channels"),
+            ]:
+                if k in cfg:
+                    load_kwargs[out_k] = cfg[k]
+        candidates = [diffusers.AutoModel]
+        if class_name:
+            maybe_cls = getattr(diffusers, str(class_name), None)
+            if maybe_cls is not None:
+                candidates.insert(0, maybe_cls)
+        last_error: Optional[Exception] = None
+        for cls in candidates:
+            try:
+                return cls.from_pretrained(comp_path, **load_kwargs)
+            except Exception as exc:
+                last_error = exc
+        raise RuntimeError(f"Failed to load transformer from {comp_path}: {last_error}")
+
+    if component == ComponentType.TEXT_ENCODER:
+        config = AutoConfig.from_pretrained(comp_path, trust_remote_code=True)
+        kwargs = {
+            "torch_dtype": torch.bfloat16,
+            "trust_remote_code": True,
+            "config": config,
+        }
+        architectures = tuple(getattr(config, "architectures", ()) or ())
+        if (
+            "UMT5EncoderModel" in architectures
+            or getattr(config, "model_type", None) == "umt5"
+        ):
+            class_order = [UMT5EncoderModel, AutoModel, AutoModelForCausalLM]
+        else:
+            class_order = [
+                AutoModel,
+                AutoModelForCausalLM,
+                UMT5EncoderModel,
+                T5EncoderModel,
+            ]
+        if AutoVisionTextModel is not None:
+            class_order.append(AutoVisionTextModel)
+        last_error: Optional[Exception] = None
+        for cls in class_order:
+            try:
+                return cls.from_pretrained(comp_path, **kwargs)
+            except Exception as exc:
+                last_error = exc
+        raise RuntimeError(
+            f"Failed to load text encoder from {comp_path}: {last_error}"
+        )
+
+    raise RuntimeError(f"Unsupported component {component.value}")
+
+
+# Public accuracy engine
+class AccuracyEngine:
+    @staticmethod
+    def reset_parallel_runtime() -> None:
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+        if model_parallel_is_initialized():
+            destroy_model_parallel()
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    @staticmethod
+    def clear_memory() -> None:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @staticmethod
+    def _execute_with_native_hook(call) -> Any:
+        output: Any = None
+
+        def _hook(_: nn.Module, __: tuple[Any, ...], captured: Any) -> None:
+            nonlocal output
+            output = captured
+
+        handle = call.module.register_forward_hook(_hook)
+        try:
+            call.module(*call.args, **call.kwargs)
+        finally:
+            handle.remove()
+        assert output is not None
+        return output
+
+    @staticmethod
+    def _apply_output_transforms(tensor: torch.Tensor, call) -> torch.Tensor:
+        if call.negate_output:
+            return -tensor
+        return tensor
+
+    @staticmethod
+    def check_accuracy(
+        target: torch.Tensor, reference: torch.Tensor, name: str, threshold: float
+    ) -> None:
+        full_tensor = getattr(target, "full_tensor", None)
+        if callable(full_tensor):
+            target = full_tensor()
+        t, r = target.detach().cpu().float(), reference.detach().cpu().float()
+
+        logger.info(
+            "[%s] Shape: SGL=%s, REF=%s | NaNs: SGL=%s, REF=%s",
+            name,
+            list(t.shape),
+            list(r.shape),
+            torch.isnan(t).sum(),
+            torch.isnan(r).sum(),
+        )
+
+        if t.shape != r.shape:
+            if t.ndim == 5 and t.shape[2] == 1:
+                t = t.squeeze(2)
+            if r.ndim == 5 and r.shape[2] == 1:
+                r = r.squeeze(2)
+            if t.shape != r.shape:
+                raise RuntimeError(
+                    f"Accuracy shape mismatch for {name}: {list(t.shape)} vs {list(r.shape)}"
+                )
+
+        cos_sim = torch.nn.functional.cosine_similarity(
+            t.reshape(-1), r.reshape(-1), dim=0
+        ).item()
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        logger.info("[%s] Rank %s CosSim=%.6f", name, rank, cos_sim)
+        assert (
+            cos_sim > threshold
+        ), f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}"
+
+    @staticmethod
+    def transfer_weights(
+        source: nn.Module,
+        target: nn.Module,
+        min_match_ratio: float = MIN_MATCH_RATIO,
+        target_device: Optional[torch.device] = None,
+    ) -> None:
+        device = target_device or get_local_torch_device()
+        dtype = torch.bfloat16
+        materialize_module(target, device, dtype)
+
+        source_state = source.state_dict()
+        mapping = getattr(target, "param_names_mapping", None) or getattr(
+            getattr(target, "module", None), "param_names_mapping", None
+        )
+        if mapping:
+            source_state, _ = hf_to_custom_state_dict(
+                source_state, get_param_names_mapping(mapping)
+            )
+
+        lookup = build_state_lookup(source_state)
+        reverse_mapping = getattr(
+            target, "reverse_param_names_mapping", None
+        ) or getattr(
+            getattr(target, "module", None), "reverse_param_names_mapping", None
+        )
+        tp_world = (
+            get_tensor_model_parallel_world_size()
+            if model_parallel_is_initialized()
+            else 1
+        )
+        rank = (
+            get_tensor_model_parallel_rank() if model_parallel_is_initialized() else 0
+        )
+        shard_contexts = build_parameter_shard_contexts(target)
+
+        matched = 0
+        total = 0
+        unmatched_details: List[str] = []
+        for name, tensor in target.named_parameters():
+            total += 1
+            src_tensor = None
+            for cand in generate_name_candidates(name, reverse_mapping):
+                if cand in lookup:
+                    src_tensor = lookup[cand]
+                    break
+            if src_tensor is None:
+                for cand in generate_name_candidates(name, reverse_mapping):
+                    src_tensor = fuse_qkv(lookup, cand)
+                    if src_tensor is not None:
+                        break
+            if src_tensor is None:
+                for cand in generate_name_candidates(name, reverse_mapping):
+                    src_tensor = fuse_gate_up_proj(lookup, cand)
+                    if src_tensor is not None:
+                        break
+            if src_tensor is None:
+                unmatched_details.append(f"{name}: no matching source tensor")
+                continue
+            shard_context = shard_contexts.get(name)
+            shard_world_size = (
+                shard_context.world_size if shard_context is not None else tp_world
+            )
+            shard_rank = shard_context.rank if shard_context is not None else rank
+            if copy_tensor(tensor, src_tensor, shard_world_size, shard_rank):
+                matched += 1
+            else:
+                unmatched_details.append(
+                    f"{name}: source {list(src_tensor.shape)} -> target {list(tensor.shape)} unsupported for shard_world_size={shard_world_size}"
+                )
+
+        for name, tensor in target.named_buffers():
+            src_tensor = None
+            for cand in generate_name_candidates(name, reverse_mapping):
+                if cand in lookup:
+                    src_tensor = lookup[cand]
+                    break
+            if src_tensor is None:
+                continue
+            shard_context = shard_contexts.get(name)
+            shard_world_size = (
+                shard_context.world_size if shard_context is not None else tp_world
+            )
+            shard_rank = shard_context.rank if shard_context is not None else rank
+            copy_tensor(tensor, src_tensor, shard_world_size, shard_rank)
+
+        ratio = matched / max(total, 1)
+        logger.info(
+            "Weight transfer: %s/%s matched (%.2f%%).", matched, total, ratio * 100
+        )
+        if ratio < min_match_ratio:
+            if rank == 0 and unmatched_details:
+                logger.error(
+                    "Unmatched parameter details:\n%s", "\n".join(unmatched_details)
+                )
+            raise RuntimeError(
+                f"Weight transfer matched {matched}/{total} ({ratio:.2%}); below threshold {min_match_ratio:.2%}."
+            )
+
+    @staticmethod
+    def run_component_pair_native(
+        case: DiffusionTestCase,
+        component: ComponentType,
+        sgl_model: nn.Module,
+        ref_model: nn.Module,
+        device: str,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if component == ComponentType.TEXT_ENCODER:
+            raise ValueError("Text encoder path is not migrated to native hooks yet")
+        profile = resolve_component_native_profile(component)
+
+        inputs = profile.build_inputs(case, sgl_model, device, ref_model)
+        sgl_call = profile.prepare_sglang_call(sgl_model, inputs)
+        ref_call = profile.prepare_reference_call(ref_model, inputs)
+
+        with torch.no_grad():
+            sgl_raw = AccuracyEngine._execute_with_native_hook(sgl_call)
+            ref_raw = AccuracyEngine._execute_with_native_hook(ref_call)
+
+        sgl_out = profile.normalize_sglang_output(sgl_raw)
+        ref_out = profile.normalize_reference_output(ref_raw)
+        sgl_out = AccuracyEngine._apply_output_transforms(sgl_out, sgl_call)
+        ref_out = AccuracyEngine._apply_output_transforms(ref_out, ref_call)
+        return sgl_out, ref_out
+
+    @staticmethod
+    def load_component_pair(
+        case: DiffusionTestCase,
+        component: ComponentType,
+        library: str,
+        num_gpus: int,
+        materialize_sgl_on_device: bool = True,
+        materialize_ref_on_device: bool = True,
+    ) -> Tuple[nn.Module, nn.Module, str]:
+        spec = COMPONENT_SPECS[component]
+        if library != spec.reference_library:
+            logger.warning(
+                "Overriding library '%s' with '%s' for component '%s'.",
+                library,
+                spec.reference_library,
+                component.value,
+            )
+            library = spec.reference_library
+        hub_id = case.server_args.model_path
+        component_selection = select_component_source(
+            hub_id,
+            case.server_args.extras,
+            component,
+            spec.model_index_keys,
+        )
+        sgl_args = build_accuracy_server_args(
+            component_selection.base_model_id,
+            component_selection.base_model_root,
+            case,
+            component,
+            num_gpus,
+            component_selection.component_paths,
+        )
+        initialize_parallel_runtime(sgl_args)
+        set_global_server_args(sgl_args)
+
+        device = get_local_torch_device()
+
+        sgl_component = _load_sglang_component(
+            component_selection.source_path,
+            sgl_args,
+            component,
+            library,
+            text_encoder_cpu_offload=(
+                False
+                if component != ComponentType.TEXT_ENCODER or materialize_sgl_on_device
+                else True
+            ),
+        )
+        if materialize_sgl_on_device:
+            sgl_component = sgl_component.to(device=device, dtype=torch.bfloat16)
+
+        ref_component = _load_reference_component(
+            component_selection.source_path,
+            component_selection.source_root,
+            component,
+            hub_id,
+            sgl_args.pipeline_config,
+            component_selection.source_subfolder,
+        )
+        if materialize_ref_on_device:
+            ref_component = ref_component.to(device=device, dtype=torch.bfloat16)
+
+        if component == ComponentType.TRANSFORMER and "wan" in hub_id.lower():
+            fc_mod._forward_context = ForwardContext(
+                current_timestep=0, attn_metadata=None
+            )
+
+        ref_for_transfer = ref_component
+        if (
+            component == ComponentType.TEXT_ENCODER
+            and getattr(ref_component, "shared", None) is None
+        ):
+            ref_for_transfer = resolve_text_encoder_forward_module(ref_component)
+        AccuracyEngine.transfer_weights(
+            ref_for_transfer,
+            sgl_component,
+            target_device=(
+                device if materialize_sgl_on_device else torch.device("cpu")
+            ),
+        )
+
+        if component != ComponentType.VAE:
+            if not hasattr(fc_mod._forward_context, "attn_metadata"):
+                fc_mod._forward_context = ForwardContext(
+                    current_timestep=int(DEFAULT_TIMESTEP), attn_metadata=None
+                )
+
+        return sgl_component.eval(), ref_component.eval(), str(device)

python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py

@@ -0,0 +1,37 @@
+import pytest
+
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    ComponentType,
+    get_skip_reason,
+    should_skip_component,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    run_native_component_accuracy_case,
+    run_text_encoder_accuracy_case,
+)
+from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
+from sglang.multimodal_gen.test.server.testcase_configs import ONE_GPU_CASES_A
+
+
+@pytest.mark.parametrize("case", ONE_GPU_CASES_A, ids=lambda x: x.id)
+class TestAccuracy1GPU_A:
+    """1-GPU Component Accuracy Suite (Set A)."""
+
+    def test_vae_accuracy(self, case):
+        if should_skip_component(case, ComponentType.VAE):
+            pytest.skip(get_skip_reason(case, ComponentType.VAE))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.VAE, "diffusers", 1
+        )
+
+    def test_transformer_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TRANSFORMER):
+            pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 1
+        )
+
+    def test_encoder_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TEXT_ENCODER):
+            pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
+        run_text_encoder_accuracy_case(AccuracyEngine, case, 1)

python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py

@@ -0,0 +1,37 @@
+import pytest
+
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    ComponentType,
+    get_skip_reason,
+    should_skip_component,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    run_native_component_accuracy_case,
+    run_text_encoder_accuracy_case,
+)
+from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
+from sglang.multimodal_gen.test.server.testcase_configs import ONE_GPU_CASES_B
+
+
+@pytest.mark.parametrize("case", ONE_GPU_CASES_B, ids=lambda x: x.id)
+class TestAccuracy1GPU_B:
+    """1-GPU Component Accuracy Suite (Set B)."""
+
+    def test_vae_accuracy(self, case):
+        if should_skip_component(case, ComponentType.VAE):
+            pytest.skip(get_skip_reason(case, ComponentType.VAE))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.VAE, "diffusers", 1
+        )
+
+    def test_transformer_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TRANSFORMER):
+            pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 1
+        )
+
+    def test_encoder_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TEXT_ENCODER):
+            pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
+        run_text_encoder_accuracy_case(AccuracyEngine, case, 1)

python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py

@@ -0,0 +1,37 @@
+import pytest
+
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    ComponentType,
+    get_skip_reason,
+    should_skip_component,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    run_native_component_accuracy_case,
+    run_text_encoder_accuracy_case,
+)
+from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
+from sglang.multimodal_gen.test.server.testcase_configs import TWO_GPU_CASES_A
+
+
+@pytest.mark.parametrize("case", TWO_GPU_CASES_A, ids=lambda x: x.id)
+class TestAccuracy2GPU_A:
+    """2-GPU Component Accuracy Suite (Set A)."""
+
+    def test_vae_accuracy(self, case):
+        if should_skip_component(case, ComponentType.VAE):
+            pytest.skip(get_skip_reason(case, ComponentType.VAE))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.VAE, "diffusers", 2
+        )
+
+    def test_transformer_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TRANSFORMER):
+            pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 2
+        )
+
+    def test_encoder_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TEXT_ENCODER):
+            pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
+        run_text_encoder_accuracy_case(AccuracyEngine, case, 2)

python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py

@@ -0,0 +1,37 @@
+import pytest
+
+from sglang.multimodal_gen.test.server.accuracy_config import (
+    ComponentType,
+    get_skip_reason,
+    should_skip_component,
+)
+from sglang.multimodal_gen.test.server.accuracy_utils import (
+    run_native_component_accuracy_case,
+    run_text_encoder_accuracy_case,
+)
+from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
+from sglang.multimodal_gen.test.server.testcase_configs import TWO_GPU_CASES_B
+
+
+@pytest.mark.parametrize("case", TWO_GPU_CASES_B, ids=lambda x: x.id)
+class TestAccuracy2GPU_B:
+    """2-GPU Component Accuracy Suite (Set B)."""
+
+    def test_vae_accuracy(self, case):
+        if should_skip_component(case, ComponentType.VAE):
+            pytest.skip(get_skip_reason(case, ComponentType.VAE))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.VAE, "diffusers", 2
+        )
+
+    def test_transformer_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TRANSFORMER):
+            pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
+        run_native_component_accuracy_case(
+            AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 2
+        )
+
+    def test_encoder_accuracy(self, case):
+        if should_skip_component(case, ComponentType.TEXT_ENCODER):
+            pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
+        run_text_encoder_accuracy_case(AccuracyEngine, case, 2)

scripts/ci/utils/slash_command_handler.py

@@ -271,6 +271,8 @@ def handle_rerun_stage(
         "stage-c-test-deepep-8-gpu-h200",
         "multimodal-gen-test-1-gpu",
         "multimodal-gen-test-2-gpu",
+        "multimodal-gen-component-accuracy-1-gpu",
+        "multimodal-gen-component-accuracy-2-gpu",
     ]
 
     # Valid AMD stage names that support target_stage