refactor(inputs): build the Omni renderer as a subclass instead of a proxy

OmniRenderer wrapped the upstream renderer, forwarded unknown attributes via
__getattr__, and re-bound eight BaseRenderer methods (render_cmpl, render_chat,
process_for_engine, _process_enc_dec and their async twins) onto the wrapper
so that they would reach Omni's _process_singleton. That allow-list hid any
concrete-renderer override of those methods and would silently skip the
Omni extras for any future BaseRenderer entry point not on the list.

Make OmniRenderer a mixin and create the renderer as a real subclass of the
upstream class instead:

- omni_renderer_cls(base_cls) builds (and caches) Omni<Base> = (OmniRenderer,
  base_cls); only _process_singleton / _process_singleton_async are extended,
  everything else is the upstream renderer unchanged.
- build_omni_renderer(vllm_config) resolves the class the way upstream
  renderer_from_config does (cached_tokenizer_from_config,
  tokenizer_args_from_config, RENDERER_REGISTRY.load_renderer_cls) and
  instantiates the Omni subclass; renderer_cls/tokenizer can be passed to
  skip the resolution.
- build_stage0_input_processor passes the renderer to InputProcessor up
  front instead of replacing it afterwards; the skip-tokenizer branch keeps
  _build_token_only_renderer and never resolves a tokenizer.

Rebinding an existing instance's __class__ was rejected: CPython refuses it
for HfRenderer ("object layout differs"), so the subclass is constructed.

Tests: tests/inputs/test_preprocess.py builds the renderer with
omni_renderer_cls and adds checks that the subclass is created once, that a
concrete renderer's render_cmpl override is no longer suppressed, and that
build_omni_renderer resolves like upstream or honours an explicit class;
tests/engine/test_async_omni_engine_stage_init.py asserts the stage-0
processor receives the built subclass instance and that the skip-tokenizer
branch never calls build_omni_renderer. Two pre-existing mypy complaints in
test_preprocess.py are fixed with annotations.

Validation: ruff, pre-commit (incl. mypy) on the four files; the modules
import and omni_renderer_cls(HfRenderer) builds; tests/inputs (32),
tests/engine/test_async_omni_engine_stage_init.py (55) and tests/entrypoints
(548) pass on CPU.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpXe2LeRrUQGjZfKFaocMZ
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/engine/test_async_omni_engine_stage_init.py

@@ -1593,22 +1593,34 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
     assert "model_path_resolver" not in engine_args
 
 
-def test_build_stage0_input_processor_wraps_active_renderer(monkeypatch):
-    import vllm_omni.engine.stage_init_utils as init_mod
-    from vllm_omni.inputs.preprocess import OmniRenderer
+def test_build_stage0_input_processor_uses_omni_renderer_subclass(monkeypatch):
+    from vllm.renderers import BaseRenderer
 
-    original = object()
+    import vllm_omni.engine.stage_init_utils as init_mod
+    from vllm_omni.inputs.preprocess import OmniRenderer, omni_renderer_cls
+
+    class _Base(BaseRenderer):
+        def __init__(self, config, tokenizer):
+            self.config, self.tokenizer = config, tokenizer
+
+        def render_messages(self, messages, params):  # pragma: no cover - abstract stub
+            raise NotImplementedError
+
+    config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
+    built = omni_renderer_cls(_Base)(config, "tok")
+    seen = {}
 
     class DummyInputProcessor:
         def __init__(self, vllm_config, renderer=None):
-            self.renderer = renderer or original
+            seen["renderer"] = renderer
+            self.renderer = renderer
 
     monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
-    processor = build_stage0_input_processor(
-        types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
-    )
+    monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None)
+    processor = build_stage0_input_processor(config)
+    assert seen["renderer"] is built
     assert isinstance(processor.renderer, OmniRenderer)
-    assert processor.renderer._renderer is original
+    assert isinstance(processor.renderer, _Base)
     assert not hasattr(processor, "input_preprocessor")
 
 
@@ -1622,13 +1634,17 @@ def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(mo
             assert renderer is original
             self.renderer = renderer
 
+    def _must_not_resolve(_cfg):
+        raise AssertionError("tokenizer must not be resolved when skip_tokenizer_init=True")
+
     monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
     monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original)
+    monkeypatch.setattr(init_mod, "build_omni_renderer", _must_not_resolve)
     config = types.SimpleNamespace(
         model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {})
     )
     processor = build_stage0_input_processor(config)
-    assert processor.renderer._renderer is original
+    assert processor.renderer is original
 
 
 def test_inject_kv_stage_info_infers_sender_tp_topology():

tests/inputs/test_preprocess.py

@@ -14,7 +14,8 @@ from vllm.renderers import BaseRenderer
 from vllm.renderers.params import TokenizeParams
 from vllm.v1.engine.input_processor import InputProcessor
 
-from vllm_omni.inputs.preprocess import OmniRenderer
+from vllm_omni.inputs import preprocess as preprocess_mod
+from vllm_omni.inputs.preprocess import OmniRenderer, build_omni_renderer, omni_renderer_cls
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -35,7 +36,7 @@ class _Renderer(BaseRenderer):
 
 @pytest.fixture
 def renderer():
-    return OmniRenderer(_Renderer())
+    return omni_renderer_cls(_Renderer)()
 
 
 @pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}])
@@ -55,12 +56,12 @@ def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, token
     processor._validate_params = Mock()
     processor._validate_lora = Mock()
     processor._validate_model_inputs = Mock()
-    prompt = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"}
+    prompt: dict[str, object] = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"}
     prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt")
     request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",))
     assert request.prompt_token_ids == [1, 2, 3, 99]
     assert request.cache_salt == "salt"
-    renderer._renderer._process_multimodal.assert_called_once_with(
+    renderer._process_multimodal.assert_called_once_with(
         [1, 2, 3],
         {},
         mm_processor_kwargs=kwargs,
@@ -72,7 +73,11 @@ def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, token
 @pytest.mark.parametrize("async_mode", [False, True])
 @pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"])
 def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind):
-    prompt = {"prompt": "hello", "additional_information": {"speaker": 1}, "model_intermediate_buffer": {"x": 2}}
+    prompt: dict[str, object] = {
+        "prompt": "hello",
+        "additional_information": {"speaker": 1},
+        "model_intermediate_buffer": {"x": 2},
+    }
     if kind == "tokens":
         prompt["prompt_token_ids"] = [1, 2, 3]
     elif kind == "media":
@@ -90,10 +95,63 @@ def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode,
     assert result["prompt"] == "hello"
     if kind in ("media", "kwargs"):
         assert result["prompt_token_ids"] == [1, 2, 3, 99]
-        assert renderer._renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True
+        assert renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True
     else:
-        renderer._renderer._process_multimodal.assert_not_called()
+        renderer._process_multimodal.assert_not_called()
         if kind == "embeds":
             torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"])
         else:
             assert result["prompt_token_ids"] == [1, 2, 3]
+
+
+def test_omni_renderer_is_a_real_subclass_created_once():
+    cls = omni_renderer_cls(_Renderer)
+    assert issubclass(cls, OmniRenderer) and issubclass(cls, _Renderer)
+    assert omni_renderer_cls(_Renderer) is cls
+    assert omni_renderer_cls(cls) is cls
+    assert cls.__name__ == "Omni_Renderer"
+    instance = cls()
+    assert isinstance(instance, _Renderer)
+    assert not hasattr(instance, "_renderer")
+
+
+def test_concrete_renderer_overrides_are_not_suppressed():
+    """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides."""
+
+    class _Custom(_Renderer):
+        def render_cmpl(self, prompts, *args, **kwargs):
+            return ["custom"]
+
+    assert omni_renderer_cls(_Custom)().render_cmpl([{"prompt": "x"}]) == ["custom"]
+
+
+class _CtorRenderer(BaseRenderer):
+    def __init__(self, config, tokenizer):
+        self.config = config
+        self.tokenizer = tokenizer
+
+    def render_messages(self, messages, params):
+        raise NotImplementedError
+
+
+def test_build_omni_renderer_with_explicit_class_skips_registry(monkeypatch):
+    monkeypatch.setattr(preprocess_mod, "cached_tokenizer_from_config", Mock(side_effect=AssertionError("resolved")))
+    config = SimpleNamespace(model_config=SimpleNamespace())
+    built = build_omni_renderer(config, renderer_cls=_CtorRenderer, tokenizer=None)
+    assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer)
+    assert built.config is config and built.tokenizer is None
+
+
+def test_build_omni_renderer_resolves_like_upstream(monkeypatch):
+    config = SimpleNamespace(model_config=SimpleNamespace())
+    tokenizer = object()
+    monkeypatch.setattr(
+        preprocess_mod, "cached_tokenizer_from_config", lambda mc: tokenizer if mc is config.model_config else None
+    )
+    monkeypatch.setattr(preprocess_mod, "tokenizer_args_from_config", lambda mc: ("fake-mode", None))
+    load_cls = Mock(return_value=_CtorRenderer)
+    monkeypatch.setattr(preprocess_mod.RENDERER_REGISTRY, "load_renderer_cls", load_cls)
+    built = build_omni_renderer(config)
+    load_cls.assert_called_once_with("fake-mode")
+    assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer)
+    assert built.tokenizer is tokenizer

vllm_omni/engine/stage_init_utils.py

@@ -54,7 +54,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
-from vllm_omni.inputs.preprocess import OmniRenderer
+from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.inc_config import OmniINCConfig
@@ -1549,22 +1549,24 @@ class _TokenOnlyRenderer(BaseRenderer):
 
 
 def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer:
-    return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None)
+    return omni_renderer_cls(_TokenOnlyRenderer)(stage_vllm_config, tokenizer=None)
 
 
 def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
-    """Build the shared stage-0 input processor."""
+    """Build the shared stage-0 input processor.
+
+    The renderer is the Omni subclass of the upstream renderer class that
+    ``renderer_from_config`` would have picked (or of the token-only renderer
+    when the stage skips tokenizer initialization), so upstream's
+    ``InputProcessor`` runs unmodified on top of it.
+    """
 
     patch_generation_config_if_needed(stage_vllm_config.model_config)
     if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)):
-        input_processor = InputProcessor(
-            vllm_config=stage_vllm_config,
-            renderer=_build_token_only_renderer(stage_vllm_config),
-        )
+        renderer = _build_token_only_renderer(stage_vllm_config)
     else:
-        input_processor = InputProcessor(vllm_config=stage_vllm_config)
-    input_processor.renderer = OmniRenderer(input_processor.renderer)
-    return input_processor
+        renderer = build_omni_renderer(stage_vllm_config)
+    return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer)
 
 
 def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str:

vllm_omni/inputs/preprocess.py

@@ -1,34 +1,46 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Omni renderer: the upstream renderer class with Omni engine-input routing.
 
-from typing import Any
+``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a
+concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream
+renderer (tokenizer, executors, multimodal cache and every public entry point
+come from upstream unchanged) and only ``_process_singleton`` /
+``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves
+the concrete class the same way upstream ``renderer_from_config`` does.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
 
 from vllm.renderers import BaseRenderer
+from vllm.renderers.registry import RENDERER_REGISTRY
+from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.tokenizers import TokenizerLike
 
 
 class OmniRenderer:
-    """Delegate tokenization/chat rendering while extending engine-input routing.
+    """Mixin that routes Omni prompts through the upstream renderer.
 
-    Keep the original renderer's tokenizer, executors, and multimodal caches.
-    The upstream pipeline methods run on this adapter so both sync and async
-    entrypoints reach the Omni singleton processing below.
+    It has no base class of its own; ``omni_renderer_cls`` prepends it to a
+    concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that
+    renderer at runtime (hence the casts below).
+
+    Two extensions over upstream ``_process_singleton``:
+
+    * a token prompt that carries ``mm_processor_kwargs`` but no
+      ``multi_modal_data`` still goes through the multimodal processor (AR
+      image-generation models such as GLM-Image read their target size from
+      those kwargs);
+    * Omni-only prompt keys (``prompt``, ``cache_salt``,
+      ``additional_information``, ``model_intermediate_buffer``) are copied
+      onto the engine input.
     """
 
-    render_cmpl = BaseRenderer.render_cmpl
-    render_cmpl_async = BaseRenderer.render_cmpl_async
-    render_chat = BaseRenderer.render_chat
-    render_chat_async = BaseRenderer.render_chat_async
-    process_for_engine = BaseRenderer.process_for_engine
-    process_for_engine_async = BaseRenderer.process_for_engine_async
-    _process_enc_dec = BaseRenderer._process_enc_dec
-    _process_enc_dec_async = BaseRenderer._process_enc_dec_async
-
-    def __init__(self, renderer: BaseRenderer) -> None:
-        self._renderer = renderer
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._renderer, name)
-
     @staticmethod
     def _with_omni_extras(inputs, prompt):
         for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"):
@@ -36,9 +48,14 @@ class OmniRenderer:
                 inputs[key] = prompt[key]
         return inputs
 
+    @staticmethod
+    def _routes_no_media_kwargs(prompt) -> bool:
+        return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data")
+
     def _process_singleton(self, prompt, *, skip_mm_cache: bool = False):
-        if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"):
-            inputs = self._renderer._process_multimodal(
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = renderer._process_multimodal(
                 prompt["prompt_token_ids"],
                 {},
                 mm_processor_kwargs=prompt["mm_processor_kwargs"],
@@ -46,12 +63,13 @@ class OmniRenderer:
                 skip_mm_cache=skip_mm_cache,
             )
         else:
-            inputs = self._renderer._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
         return self._with_omni_extras(inputs, prompt)
 
     async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False):
-        if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"):
-            inputs = await self._renderer._process_multimodal_async(
+        renderer = cast(BaseRenderer, self)
+        if self._routes_no_media_kwargs(prompt):
+            inputs = await renderer._process_multimodal_async(
                 prompt["prompt_token_ids"],
                 {},
                 mm_processor_kwargs=prompt["mm_processor_kwargs"],
@@ -59,5 +77,42 @@ class OmniRenderer:
                 skip_mm_cache=skip_mm_cache,
             )
         else:
-            inputs = await self._renderer._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
         return self._with_omni_extras(inputs, prompt)
+
+
+_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {}
+
+
+def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]:
+    """Return the Omni subclass of ``base_cls`` (created once per base class)."""
+    if issubclass(base_cls, OmniRenderer):
+        return base_cls
+    cls = _OMNI_RENDERER_CLASSES.get(base_cls)
+    if cls is None:
+        name = f"Omni{base_cls.__name__}"
+        cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name})
+        _OMNI_RENDERER_CLASSES[base_cls] = cls
+    return cls
+
+
+def build_omni_renderer(
+    vllm_config: VllmConfig,
+    *,
+    renderer_cls: type[BaseRenderer] | None = None,
+    tokenizer: TokenizerLike | None = None,
+) -> BaseRenderer:
+    """Build the Omni renderer for ``vllm_config``.
+
+    With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``:
+    the tokenizer and renderer mode come from the model config and the class
+    from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally
+    ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages.
+    """
+    model_config = vllm_config.model_config
+    if renderer_cls is None:
+        if tokenizer is None:
+            tokenizer = cached_tokenizer_from_config(model_config)
+        renderer_mode, *_ = tokenizer_args_from_config(model_config)
+        renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode)
+    return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer)