[Perf] Deduplicate AR prefix cache hidden-state CPU staging (#3734)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

vllm_omni/entrypoints/openai/api_server.py

@@ -17,6 +17,7 @@ from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
 from numbers import Integral
+from pathlib import Path
 from typing import Annotated, Any, Literal, cast
 
 import httpx
@@ -143,6 +144,51 @@ MAX_UINT32_SEED = 2**32 - 1
 profiler_router = APIRouter()
 
 
+def _load_model_chat_template_json(model: str) -> str | None:
+    """Load a model-level chat_template.json from a local path or HF cache.
+
+    Some multimodal HF repos, including Qwen3-Omni, ship the chat template as a
+    separate file instead of embedding it in tokenizer_config.json. Transformers
+    4.44+ no longer supplies a default template, so serving must pass that model
+    template explicitly when the user did not provide --chat-template.
+    """
+    candidate = Path(model) / "chat_template.json"
+    template_path: str | None = str(candidate) if candidate.is_file() else None
+
+    if template_path is None:
+        try:
+            from huggingface_hub import hf_hub_download
+
+            template_path = hf_hub_download(
+                repo_id=model,
+                filename="chat_template.json",
+                local_files_only=True,
+            )
+        except Exception:
+            return None
+
+    try:
+        with open(template_path, encoding="utf-8") as f:
+            payload = json.load(f)
+    except Exception as exc:
+        logger.warning("Failed to load chat template from %s: %s", template_path, exc)
+        return None
+
+    if isinstance(payload, dict):
+        template = payload.get("chat_template")
+    elif isinstance(payload, str):
+        template = payload
+    else:
+        template = None
+
+    if not isinstance(template, str) or not template.strip():
+        logger.warning("Ignoring malformed chat template payload in %s", template_path)
+        return None
+
+    logger.info("Loaded chat template from %s", template_path)
+    return template
+
+
 def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool:
     """Check if any stage has profiler_config set in its engine_args."""
     if not stage_configs:
@@ -673,6 +719,14 @@ async def omni_init_app_state(
     logger.info("Supported tasks: %s", supported_tasks)
 
     resolved_chat_template = load_chat_template(args.chat_template)
+    if resolved_chat_template is None:
+        try:
+            tokenizer = await engine_client.get_tokenizer()
+        except Exception as exc:
+            logger.debug("Could not inspect tokenizer chat_template before serving init: %s", exc)
+            tokenizer = None
+        if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
+            resolved_chat_template = _load_model_chat_template_json(args.model)
 
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()