[diffusion][model] Add LingBot Video dense and MoE support (#5035)

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Signed-off-by: wtz2333 <2955110911@qq.com>
Co-authored-by: hsliuustc0106 <liuhongsheng4@huawei.com>

benchmarks/README.md

@@ -28,6 +28,14 @@ Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offl
 - **Tasks**: text-to-image and image-to-image
 - **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
 
+### [LingBot-Video](lingbot_video/README.md) — Dense and MoE Parity
+
+Manual cross-runtime validation for the dense LingBot-Video pipeline and the
+LingBot-Video MoE transformer.
+
+- **Dense pipeline**: decoded-video MAE, MSE, PSNR, latency, and optional steady-state timings
+- **MoE transformer**: bitwise router, sparse-block, shared-expert, and full-transformer parity
+
 ### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
 
 RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.

benchmarks/lingbot_video/README.md

@@ -0,0 +1,62 @@
+# LingBot-Video Benchmarks
+
+These manual benchmarks compare the upstream LingBot-Video implementation with
+vLLM-Omni. Run all commands from the repository root.
+
+| Benchmark | Script | Purpose |
+|-----------|--------|---------|
+| Dense pipeline parity | `dense_pipeline_parity.py` | Compare generated videos and report MAE, MSE, PSNR, latency, and optional steady-state timings |
+| MoE numerical parity | `moe_transformer_parity.py` | Check bitwise parity for the sparse MoE block and the full MoE transformer |
+
+Both benchmarks require a local LingBot-Video checkout and local model files.
+They do not run in CI.
+
+## Dense pipeline parity
+
+The dense benchmark launches the upstream Diffusers inference script and the
+vLLM-Omni offline example with identical prompts and sampling parameters. It
+then compares the decoded MP4 frames.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+python benchmarks/lingbot_video/dense_pipeline_parity.py \
+  --model /path/to/lingbot-video-dense-1.3b \
+  --official-repo /path/to/lingbot-video \
+  --output-dir /tmp/lingbot_dense_parity
+```
+
+Use `--runs N` to additionally measure repeated in-process vLLM-Omni
+requests. The video comparison reports MAE, MSE, and PSNR; it is an end-to-end
+pipeline diagnostic rather than a bitwise transformer check.
+
+## MoE numerical parity
+
+The lightweight block comparison uses deterministic weights and inputs to
+cover correction-bias routing, group-limited top-k, routed experts, FP32
+scatter-weighted restore, padding masks, and the shared expert.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+python benchmarks/lingbot_video/moe_transformer_parity.py \
+  --scope block \
+  --official-repo /path/to/lingbot-video \
+  --output-json /tmp/lingbot_moe_block_parity.json
+```
+
+The full comparison loads the upstream and vLLM-Omni 30B transformers
+sequentially and compares their output tensors. It requires a CUDA device with
+`torch._grouped_mm` and enough memory for one 30B BF16 transformer.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+python benchmarks/lingbot_video/moe_transformer_parity.py \
+  --scope transformer \
+  --official-repo /path/to/lingbot-video \
+  --model /path/to/lingbot-video-moe-30b-a3b \
+  --output-json /tmp/lingbot_moe_transformer_parity.json
+```
+
+For the correctness oracle, the upstream transformer uses
+`diffusers:_native_math` and vLLM-Omni uses
+`TORCH_SDPA + SDPBackend.MATH`. The command exits with a nonzero status unless
+the selected comparisons are bitwise equal.

benchmarks/lingbot_video/dense_pipeline_parity.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import imageio.v3 as iio
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
+    parser.add_argument("--official-repo", default="/tmp/lingbot-video")
+    parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.")
+    return parser.parse_args()
+
+
+def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]:
+    start = time.perf_counter()
+    proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    elapsed = time.perf_counter() - start
+    if proc.returncode != 0:
+        print(proc.stdout)
+        raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout)
+    return elapsed, proc.stdout
+
+
+def _read_video(path: Path) -> np.ndarray:
+    frames = iio.imread(path, plugin="pyav")
+    if frames.dtype != np.float32:
+        frames = frames.astype(np.float32) / 255.0
+    return frames
+
+
+def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
+    if a.shape != b.shape:
+        return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False}
+    diff = a - b
+    mse = float(np.mean(diff * diff))
+    mae = float(np.mean(np.abs(diff)))
+    psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
+    return {
+        "shape": list(a.shape),
+        "shape_match": True,
+        "mae": mae,
+        "mse": mse,
+        "psnr": psnr,
+    }
+
+
+def _extract_total_generation_seconds(log: str) -> float | None:
+    match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log)
+    return float(match.group(1)) if match else None
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]:
+    load_start = time.perf_counter()
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+    load_seconds = time.perf_counter() - load_start
+    prompt = {"prompt": args.prompt}
+    latencies: list[float] = []
+    peak_mb = 0.0
+
+    for _ in range(args.runs):
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            num_inference_steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            generator=generator,
+            output_type="pt",
+            extra_args={"shift": args.shift},
+        )
+        start = time.perf_counter()
+        output = omni.generate(prompt, sampling_params)
+        latencies.append(time.perf_counter() - start)
+        peak_mb = max(peak_mb, _peak_memory_mb(output))
+        del output
+    omni.close()
+
+    return {
+        "load_seconds": load_seconds,
+        "run_seconds": latencies,
+        "mean_seconds": statistics.mean(latencies) if latencies else None,
+        "median_seconds": statistics.median(latencies) if latencies else None,
+        "min_seconds": min(latencies) if latencies else None,
+        "max_seconds": max(latencies) if latencies else None,
+        "peak_mb": peak_mb,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    official_out = output_dir / "official_diffusers.mp4"
+    native_out = output_dir / "vllm_omni_native.mp4"
+
+    official_cmd = [
+        sys.executable,
+        str(Path(args.official_repo) / "scripts" / "inference.py"),
+        "--backend",
+        "diffusers",
+        "--model_dir",
+        args.model,
+        "--mode",
+        "t2v",
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(official_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num_frames",
+        str(args.num_frames),
+        "--steps",
+        str(args.steps),
+        "--guidance_scale",
+        str(args.guidance_scale),
+        "--shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+        "--quiet_progress",
+    ]
+    native_cmd = [
+        sys.executable,
+        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
+        "--model",
+        args.model,
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(native_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num-frames",
+        str(args.num_frames),
+        "--num-inference-steps",
+        str(args.steps),
+        "--guidance-scale",
+        str(args.guidance_scale),
+        "--flow-shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+    ]
+
+    official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo))
+    native_time, native_log = _run(native_cmd, cwd=Path.cwd())
+
+    result = {
+        "official_seconds": official_time,
+        "native_process_seconds": native_time,
+        "native_request_seconds": _extract_total_generation_seconds(native_log),
+        "official_output": str(official_out),
+        "native_output": str(native_out),
+        "accuracy": _compare(_read_video(official_out), _read_video(native_out)),
+        "official_log_tail": official_log.splitlines()[-20:],
+        "native_log_tail": native_log.splitlines()[-20:],
+    }
+    if args.runs > 0:
+        result["native_steady_state"] = _benchmark_native_requests(args)
+
+    result_path = output_dir / "comparison.json"
+    result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
+    print(json.dumps(result, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/lingbot_video/moe_transformer_parity.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Reproduce LingBot-Video MoE block and transformer numerical parity."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import importlib
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--scope",
+        choices=("block", "transformer", "all"),
+        default="block",
+        help="Run the lightweight sparse-block check, the real transformer check, or both.",
+    )
+    parser.add_argument(
+        "--official-repo",
+        type=Path,
+        required=True,
+        help="Local checkout of https://github.com/Robbyant/lingbot-video.",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        help="Local MoE checkpoint path or a checkpoint already present in the Hugging Face cache.",
+    )
+    parser.add_argument("--transformer-subfolder", default="transformer")
+    parser.add_argument("--output-json", type=Path, default=None)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--num-frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=8)
+    parser.add_argument("--width", type=int, default=8)
+    parser.add_argument("--text-length", type=int, default=16)
+    return parser.parse_args()
+
+
+def _load_official_module(official_repo: Path):
+    transformer_file = official_repo / "lingbot_video" / "transformer_lingbot_video.py"
+    if not transformer_file.is_file():
+        raise FileNotFoundError(f"Upstream transformer not found: {transformer_file}")
+    sys.path.insert(0, str(official_repo))
+    module = importlib.import_module("lingbot_video.transformer_lingbot_video")
+    module_path = Path(module.__file__).resolve()
+    if official_repo.resolve() not in module_path.parents:
+        raise RuntimeError(f"Imported LingBot module from {module_path}, not {official_repo.resolve()}")
+    return module
+
+
+def _load_native_module():
+    return importlib.import_module("vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer")
+
+
+def _tensor_metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, Any]:
+    actual_float = actual.float()
+    expected_float = expected.float()
+    diff = (actual_float - expected_float).abs()
+    denominator = actual_float.norm() * expected_float.norm()
+    cosine = (
+        float(torch.dot(actual_float.flatten(), expected_float.flatten()) / denominator)
+        if float(denominator) > 0
+        else 1.0
+    )
+    return {
+        "shape": list(actual.shape),
+        "equal": bool(torch.equal(actual, expected)),
+        "max_abs": float(diff.max()) if diff.numel() else 0.0,
+        "mean_abs": float(diff.mean()) if diff.numel() else 0.0,
+        "rmse": float(torch.sqrt(torch.mean(diff.square()))) if diff.numel() else 0.0,
+        "cosine": cosine,
+        "finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()),
+    }
+
+
+def _initialize_block(module: torch.nn.Module, seed: int) -> None:
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    with torch.no_grad():
+        for parameter in module.parameters():
+            values = torch.randn(parameter.shape, generator=generator, dtype=torch.float32)
+            parameter.copy_((values * 0.02).to(parameter.dtype))
+        module.router.e_score_correction_bias.copy_(torch.tensor([0.9, 0.8, 1.0, 0.0, 0.7, 0.6, 0.5, 0.4]))
+
+
+def _apply_padding(
+    scores: torch.Tensor,
+    padding_mask: torch.Tensor,
+    route_scale: float,
+) -> torch.Tensor:
+    scores = scores * padding_mask.unsqueeze(-1).to(scores.dtype)
+    scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-9)
+    return scores * route_scale
+
+
+def _run_block_parity(
+    official_module,
+    native_module,
+    *,
+    device: torch.device,
+    seed: int,
+) -> dict[str, Any]:
+    if not hasattr(torch, "_grouped_mm"):
+        raise RuntimeError("The block parity check requires torch._grouped_mm.")
+
+    common = {
+        "hidden_size": 16,
+        "num_experts": 8,
+        "top_k": 2,
+        "moe_intermediate_size": 8,
+        "score_func": "sigmoid",
+        "norm_topk_prob": True,
+        "n_group": 4,
+        "topk_group": 1,
+        "routed_scaling_factor": 1.5,
+        "n_shared_experts": 1,
+    }
+    official = official_module.LingBotVideoSparseMoeBlock(
+        intermediate_size=32,
+        **common,
+    )
+    native = native_module.LingBotVideoSparseMoeBlock(**common)
+    _initialize_block(official, seed)
+    native.load_state_dict(official.state_dict(), strict=True)
+    official = official.to(device=device, dtype=torch.bfloat16).eval()
+    native = native.to(device=device, dtype=torch.bfloat16).eval()
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+    hidden_states = torch.randn(
+        2,
+        5,
+        common["hidden_size"],
+        generator=generator,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    padding_mask = torch.tensor(
+        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
+        device=device,
+        dtype=torch.float32,
+    )
+    tokens = hidden_states.reshape(-1, common["hidden_size"])
+
+    with torch.inference_mode():
+        official_router = official.router(tokens)
+        official_indices, official_scores = official_router[:2]
+        native_indices, native_scores = native.router(tokens)
+
+        official_scores = _apply_padding(
+            official_scores,
+            padding_mask,
+            official.router.route_scale,
+        )
+        native_scores = _apply_padding(
+            native_scores,
+            padding_mask,
+            native.router.route_scale,
+        )
+        official_routed = official._run_selected_experts(
+            tokens,
+            official_scores,
+            official_indices,
+        )
+        native_routed = native._run_selected_experts(
+            tokens,
+            native_scores,
+            native_indices,
+        )
+        official_shared = official.shared_experts(hidden_states)
+        native_shared = native.shared_experts(hidden_states)
+        official_output = official(hidden_states, padding_mask=padding_mask)
+        native_output = native(hidden_states, padding_mask=padding_mask)
+
+    result = {
+        "router_indices_equal": bool(torch.equal(native_indices, official_indices)),
+        "router_scores": _tensor_metrics(native_scores, official_scores),
+        "routed_output": _tensor_metrics(native_routed, official_routed),
+        "shared_output": _tensor_metrics(native_shared, official_shared),
+        "final_output": _tensor_metrics(native_output, official_output),
+    }
+    result["exact"] = bool(
+        result["router_indices_equal"]
+        and result["router_scores"]["equal"]
+        and result["routed_output"]["equal"]
+        and result["shared_output"]["equal"]
+        and result["final_output"]["equal"]
+    )
+    return result
+
+
+def _release_cuda() -> None:
+    gc.collect()
+    torch.accelerator.empty_cache()
+    torch.accelerator.synchronize()
+
+
+def _make_transformer_inputs(
+    config,
+    args: argparse.Namespace,
+) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    hidden_states = torch.randn(
+        1,
+        int(config.in_channels),
+        args.num_frames,
+        args.height,
+        args.width,
+        generator=generator,
+        dtype=torch.float32,
+    ).to(torch.bfloat16)
+    encoder_hidden_states = torch.randn(
+        1,
+        args.text_length,
+        int(config.text_dim),
+        generator=generator,
+        dtype=torch.float32,
+    ).to(torch.bfloat16)
+    return {
+        "hidden_states": hidden_states,
+        "timestep": torch.tensor([500.0], dtype=torch.float32),
+        "encoder_hidden_states": encoder_hidden_states,
+        "encoder_attention_mask": torch.ones(1, args.text_length, dtype=torch.long),
+    }
+
+
+def _load_transformer(
+    transformer_cls,
+    args: argparse.Namespace,
+    device: torch.device,
+):
+    start = time.perf_counter()
+    model = transformer_cls.from_pretrained(
+        args.model,
+        subfolder=args.transformer_subfolder,
+        torch_dtype=torch.bfloat16,
+        local_files_only=True,
+        low_cpu_mem_usage=True,
+    )
+    model = model.to(device=device, dtype=torch.bfloat16).eval()
+    torch.accelerator.synchronize()
+    return model, time.perf_counter() - start
+
+
+def _forward_transformer(
+    model,
+    inputs: dict[str, torch.Tensor],
+    device: torch.device,
+    attention_context,
+) -> tuple[torch.Tensor, float, float]:
+    torch.accelerator.reset_peak_memory_stats(device)
+    gpu_inputs = {name: value.to(device) for name, value in inputs.items()}
+    start = time.perf_counter()
+    with torch.inference_mode(), attention_context:
+        output = model(**gpu_inputs, return_dict=False)[0]
+    torch.accelerator.synchronize()
+    elapsed = time.perf_counter() - start
+    peak_gib = torch.accelerator.max_memory_reserved(device) / (1024**3)
+    return output.cpu(), elapsed, peak_gib
+
+
+def _run_transformer_parity(
+    official_module,
+    native_module,
+    *,
+    args: argparse.Namespace,
+    device: torch.device,
+) -> dict[str, Any]:
+    if args.model is None:
+        raise ValueError("--model is required for transformer parity.")
+
+    from diffusers.models.attention_dispatch import attention_backend
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
+
+    official, official_load = _load_transformer(
+        official_module.LingBotVideoTransformer3DModel,
+        args,
+        device,
+    )
+    inputs = _make_transformer_inputs(official.config, args)
+    official_output, official_forward, official_peak = _forward_transformer(
+        official,
+        inputs,
+        device,
+        attention_backend("_native_math"),
+    )
+    del official
+    _release_cuda()
+
+    native_config = OmniDiffusionConfig(
+        diffusion_attention_config=AttentionConfig(
+            default=AttentionSpec(backend="TORCH_SDPA"),
+        ),
+    )
+    with set_current_diffusion_config(native_config):
+        native, native_load = _load_transformer(
+            native_module.LingBotVideoTransformer3DModel,
+            args,
+            device,
+        )
+    native_backend_prefs = {block.attn.attn.backend_pref for block in native.blocks}
+    if native_backend_prefs != {"TORCH_SDPA"}:
+        raise RuntimeError(
+            f"Failed to force the native transformer to TORCH_SDPA: resolved preferences were {native_backend_prefs}."
+        )
+    native_output, native_forward, native_peak = _forward_transformer(
+        native,
+        inputs,
+        device,
+        sdpa_kernel(SDPBackend.MATH),
+    )
+    del native
+    _release_cuda()
+
+    metrics = _tensor_metrics(native_output, official_output)
+    return {
+        "backends": {
+            "official": "diffusers:_native_math",
+            "native": "TORCH_SDPA+SDPBackend.MATH",
+        },
+        "official": {
+            "load_seconds": official_load,
+            "forward_seconds": official_forward,
+            "peak_reserved_gib": official_peak,
+        },
+        "native": {
+            "load_seconds": native_load,
+            "forward_seconds": native_forward,
+            "peak_reserved_gib": native_peak,
+        },
+        "output": metrics,
+        "exact": metrics["equal"],
+    }
+
+
+def main() -> int:
+    args = parse_args()
+    device = torch.device(args.device)
+    if device.type != "cuda" or not torch.cuda.is_available():
+        raise RuntimeError("LingBot MoE parity requires an available CUDA device.")
+
+    os.environ["DIFFUSERS_ATTN_BACKEND"] = "_native_math"
+    os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
+    os.environ["LINGBOT_MOE_EXPERT_BACKEND"] = "grouped_mm"
+    os.environ["LINGBOT_MOE_PAD_BACKEND"] = "loop"
+    os.environ["LINGBOT_MOE_REORDER_BACKEND"] = "sort"
+    os.environ["LINGBOT_MOE_RESTORE_BACKEND"] = "scatter"
+
+    official_module = _load_official_module(args.official_repo)
+    native_module = _load_native_module()
+
+    result: dict[str, Any] = {
+        "settings": {
+            "scope": args.scope,
+            "seed": args.seed,
+            "device": str(device),
+            "official_repo": str(args.official_repo.resolve()),
+            "model": args.model,
+        }
+    }
+    if args.scope in {"block", "all"}:
+        result["block"] = _run_block_parity(
+            official_module,
+            native_module,
+            device=device,
+            seed=args.seed,
+        )
+    if args.scope in {"transformer", "all"}:
+        result["transformer"] = _run_transformer_parity(
+            official_module,
+            native_module,
+            args=args,
+            device=device,
+        )
+
+    result["exact"] = all(section["exact"] for name, section in result.items() if name in {"block", "transformer"})
+    payload = json.dumps(result, indent=2, sort_keys=True)
+    print(payload)
+    if args.output_json is not None:
+        args.output_json.parent.mkdir(parents=True, exist_ok=True)
+        args.output_json.write_text(payload + "\n", encoding="utf-8")
+    return 0 if result["exact"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())