[Perf] Enable fused RMSNorm for HunyuanImage3 (#3959)

benchmarks/accuracy/common.py

@@ -151,6 +151,9 @@ class VllmOmniImageClient:
         seed: int | None = None,
         negative_prompt: str | None = None,
         output_compression: int | None = None,
+        bot_task: str | None = None,
+        sys_type: str | None = None,
+        system_prompt: str | None = None,
     ) -> Image.Image:
         if not isinstance(images, list):
             images = [images]
@@ -170,6 +173,12 @@ class VllmOmniImageClient:
             data["negative_prompt"] = negative_prompt
         if output_compression is not None:
             data["output_compression"] = str(output_compression)
+        if bot_task is not None:
+            data["bot_task"] = bot_task
+        if sys_type is not None:
+            data["sys_type"] = sys_type
+        if system_prompt is not None:
+            data["system_prompt"] = system_prompt
 
         files = [
             (

benchmarks/accuracy/image_to_image/gedit_bench.py

@@ -64,6 +64,47 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
     raise ValueError("model-name is required when it cannot be inferred from model or output-root")
 
 
+def _coerce_score_int(value: Any) -> int:
+    if isinstance(value, bool):
+        raise TypeError(f"Expected numeric score, got bool: {value!r}")
+    if isinstance(value, (int, float)):
+        return int(value)
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped.lstrip("-").isdigit():
+            return int(stripped)
+    raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
+
+
+def _expand_score_entry(entry: Any) -> list[int]:
+    if isinstance(entry, dict):
+        if "score" in entry:
+            return [_coerce_score_int(entry["score"])]
+        if "value" in entry:
+            return [_coerce_score_int(entry["value"])]
+        numeric_values = [_coerce_score_int(value) for value in entry.values()]
+        if numeric_values:
+            return numeric_values
+        raise TypeError(f"Score dict has no numeric values: {entry!r}")
+    return [_coerce_score_int(entry)]
+
+
+def _normalize_score_list(score: Any) -> list[int]:
+    if isinstance(score, dict):
+        entries: list[Any] = [score]
+    elif isinstance(score, list):
+        entries = score
+    else:
+        entries = [score]
+
+    normalized: list[int] = []
+    for entry in entries:
+        normalized.extend(_expand_score_entry(entry))
+    if not normalized:
+        raise ValueError(f"Empty score list after normalization: {score!r}")
+    return normalized
+
+
 def parse_score_payload(raw_text: str) -> dict[str, Any]:
     try:
         parsed = extract_json_object(raw_text)
@@ -73,10 +114,7 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
             parsed = {"score": json.loads(stripped), "reasoning": ""}
         else:
             raise
-    score = parsed.get("score", [])
-    if not isinstance(score, list):
-        score = [score]
-    parsed["score"] = [int(value) for value in score]
+    parsed["score"] = _normalize_score_list(parsed.get("score", []))
     return parsed
 
 
@@ -391,6 +429,9 @@ class GEditBenchRunner:
         num_inference_steps: int = 20,
         guidance_scale: float | None = None,
         seed: int | None = 42,
+        bot_task: str | None = None,
+        sys_type: str | None = None,
+        system_prompt: str | None = None,
     ):
         self.dataset_ref = dataset_ref
         self.output_root = output_root
@@ -400,6 +441,9 @@ class GEditBenchRunner:
         self.num_inference_steps = num_inference_steps
         self.guidance_scale = guidance_scale
         self.seed = seed
+        self.bot_task = bot_task
+        self.sys_type = sys_type
+        self.system_prompt = system_prompt
         self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
 
     def generate(
@@ -478,6 +522,9 @@ class GEditBenchRunner:
             num_inference_steps=self.num_inference_steps,
             guidance_scale=self.guidance_scale,
             seed=self.seed,
+            bot_task=self.bot_task,
+            sys_type=self.sys_type,
+            system_prompt=self.system_prompt,
         )
         edited_image.save(output_path)
         return {
@@ -633,6 +680,25 @@ def build_parser() -> argparse.ArgumentParser:
     generate.add_argument("--num-inference-steps", type=int, default=20)
     generate.add_argument("--guidance-scale", type=float, default=None)
     generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument(
+        "--bot-task",
+        type=str,
+        default=None,
+        choices=["think", "recaption", "think_recaption", "vanilla"],
+        help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
+    )
+    generate.add_argument(
+        "--sys-type",
+        type=str,
+        default=None,
+        help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
+    )
+    generate.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help="Custom Hunyuan system prompt body (sys_type=custom).",
+    )
     generate.add_argument("--workers", type=int, default=1)
     generate.add_argument("--max-samples", type=int, default=None)
     generate.add_argument("--samples-per-group", type=int, default=None)
@@ -675,6 +741,9 @@ def main(argv: list[str] | None = None) -> int:
             num_inference_steps=args.num_inference_steps,
             guidance_scale=args.guidance_scale,
             seed=args.seed,
+            bot_task=args.bot_task,
+            sys_type=args.sys_type,
+            system_prompt=args.system_prompt,
         )
         records = runner.generate(
             model_name=model_name,

tests/benchmarks/test_accuracy_bench_utils.py

@@ -345,6 +345,14 @@ def test_parse_score_payload_handles_raw_json_and_delimited_json():
     assert parse_score_payload(wrapped)["score"] == [6]
 
 
+def test_parse_score_payload_handles_qwen_vl_nested_score_dicts():
+    qwen_style = '{"score": [{"naturalness": 8, "artifact_free": 7}], "reasoning": "good quality"}'
+    assert parse_score_payload(qwen_style)["score"] == [8, 7]
+
+    nested_values = '{"score": [{"score": 9}, {"score": 6}], "reasoning": "ok"}'
+    assert parse_score_payload(nested_values)["score"] == [9, 6]
+
+
 def test_summarize_gedit_generated_records_groups_by_task_and_language():
     records = []
     for group in GEDIT_GROUPS[:2]:

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -32,7 +32,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
@@ -73,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding

vllm_omni/entrypoints/openai/api_server.py

@@ -2016,12 +2016,12 @@ async def edit_images(
         _update_if_not_none(gen_params, "resolution", resolution)
 
         extra_args = dict(getattr(gen_params, "extra_args", {}) or {})
-        if bot_task is not None:
-            extra_args["bot_task"] = bot_task
-        if sys_type is not None:
-            extra_args["use_system_prompt"] = sys_type
-        if system_prompt is not None:
-            extra_args["system_prompt"] = system_prompt
+        edit_extra_args = _build_hunyuan_edit_extra_args(
+            bot_task=bot_task,
+            sys_type=sys_type,
+            system_prompt=system_prompt,
+        )
+        extra_args.update(edit_extra_args)
         if extra_args:
             gen_params.extra_args = extra_args
 
@@ -2323,6 +2323,28 @@ def _check_max_generated_image_size(
             )
 
 
+def _build_hunyuan_edit_extra_args(
+    *,
+    bot_task: str | None,
+    sys_type: str | None,
+    system_prompt: str | None,
+) -> dict[str, Any]:
+    """Map Hunyuan /v1/images/edits form fields to DiT ``extra_args``."""
+    extra_args: dict[str, Any] = {}
+    effective_use_system_prompt = sys_type
+    if effective_use_system_prompt is None and bot_task is not None:
+        from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import resolve_sys_type
+
+        effective_use_system_prompt = resolve_sys_type(bot_task)
+    if effective_use_system_prompt is not None:
+        extra_args["use_system_prompt"] = effective_use_system_prompt
+    if system_prompt is not None:
+        extra_args["system_prompt"] = system_prompt
+    if bot_task is not None:
+        extra_args["bot_task"] = bot_task
+    return extra_args
+
+
 def _update_if_not_none(object: Any, key: str, val: Any) -> None:
     if val is not None:
         setattr(object, key, val)