[Bugfix][MiniCPM-o] Wire the duplex sampling hook into the NPU AR model runner (#5895)

Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/worker/test_native_duplex_hooks.py

@@ -2466,3 +2466,165 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio():
 
     assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5]
     assert len(state.context_embeds) == 6
+
+
+# ---------------------------------------------------------------------------
+# Cross-platform wiring: every AR model runner must reach the duplex hook.
+#
+# ``GPUARModelRunner`` and ``NPUARModelRunner`` are siblings -- neither inherits
+# the other's ``_sample`` -- so the NPU runner once carried a copy of the GPU
+# ``_sample`` with the ``prepare_duplex_sampling`` call dropped.  That left the
+# MiniCPM-o 4.5 thinker without a row -> session map on Ascend, so every mid-turn
+# ``<|listen|>`` was emitted verbatim and the model stopped speaking after turn 1.
+#
+# These checks parse the source instead of importing it: the NPU runner needs
+# ``vllm_ascend``, which is absent on the CPU test images that run this file.
+# ---------------------------------------------------------------------------
+
+# method that must call it -> mixin method it must call
+_REQUIRED_DUPLEX_CALLS = {
+    "__init__": "_init_duplex_sampling_state",
+    "load_model": "_resolve_duplex_sampling_hook",
+    "_update_states": "_update_duplex_sampling_states",
+    "_sample": "_apply_duplex_sampling",
+}
+
+# The sampler ``_sample`` hands the prepared logits to.  ``_apply_duplex_sampling``
+# both masks those logits (force_listen) and publishes the row -> session map the
+# model's own sampler reads, so calling it *after* this is the same silent no-op as
+# not calling it at all.
+_MODEL_SAMPLER_CALL = "model_sample"
+
+
+def _repo_root():
+    from pathlib import Path
+
+    return Path(__file__).parents[2]
+
+
+def _ar_runner_sources():
+    """Every AR model runner in the repo, discovered rather than hand-listed.
+
+    A new platform runner is picked up automatically instead of silently escaping
+    the checks below.  ``vllm_omni/worker/gpu_ar_model_runner.py`` is named
+    separately because it does not live under ``platforms/``.
+    """
+    root = _repo_root()
+    paths = [root / "vllm_omni/worker/gpu_ar_model_runner.py"]
+    paths.extend(sorted(root.glob("vllm_omni/platforms/*/worker/*_ar_model_runner.py")))
+    return paths
+
+
+def _ar_runner_classdefs():
+    """``{class_name: (ClassDef, relative_path)}`` for every ``*ARModelRunner``."""
+    import ast
+
+    root = _repo_root()
+    found = {}
+    for path in _ar_runner_sources():
+        module = ast.parse(path.read_text())
+        for node in module.body:
+            if isinstance(node, ast.ClassDef) and node.name.endswith("ARModelRunner"):
+                found[node.name] = (node, str(path.relative_to(root)))
+    assert found, "no AR model runners discovered; the glob in _ar_runner_sources() is stale"
+    return found
+
+
+def _method_calls(method):
+    """``{attribute_name: first lineno}`` for every ``x.y(...)`` call in ``method``."""
+    import ast
+
+    calls = {}
+    for node in ast.walk(method):
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+            calls.setdefault(node.func.attr, node.lineno)
+    return calls
+
+
+def _wires_duplex_sampling_directly(classdef) -> bool:
+    import ast
+
+    bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+    if "DuplexSamplingRunnerMixin" not in bases:
+        return False
+    methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
+    return all(
+        method_name in methods and required_call in _method_calls(methods[method_name])
+        for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items()
+    )
+
+
+def test_every_ar_runner_reaches_the_duplex_sampling_hook():
+    """Each AR runner must wire the hook itself or inherit from one that does.
+
+    ``XPUARModelRunner`` derives from ``GPUARModelRunner`` and so is wired for free;
+    ``NPUARModelRunner`` derives from ``OmniNPUModelRunner``, which is not an AR
+    runner, so it has to do the wiring itself.  Asserting every discovered class
+    wires directly would wrongly fail the XPU runner.
+    """
+    import ast
+
+    classdefs = _ar_runner_classdefs()
+    wired = {name for name, (classdef, _path) in classdefs.items() if _wires_duplex_sampling_directly(classdef)}
+
+    for class_name, (classdef, path) in sorted(classdefs.items()):
+        if class_name in wired:
+            continue
+        base_names = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+        assert base_names & wired, (
+            f"{path}::{class_name} neither wires the duplex sampling hook nor derives from a "
+            f"runner that does (wired runners: {sorted(wired)}).  Inherit DuplexSamplingRunnerMixin "
+            f"and call {sorted(_REQUIRED_DUPLEX_CALLS.values())} from "
+            f"{sorted(_REQUIRED_DUPLEX_CALLS)} respectively."
+        )
+
+
+@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
+def test_ar_runner_calls_every_duplex_sampling_hook_site(class_name: str):
+    import ast
+
+    classdefs = _ar_runner_classdefs()
+    assert class_name in classdefs, f"{class_name} not found by _ar_runner_sources()"
+    classdef, path = classdefs[class_name]
+
+    bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+    assert "DuplexSamplingRunnerMixin" in bases, (
+        f"{path}::{class_name} must inherit DuplexSamplingRunnerMixin so the duplex "
+        f"sampling hook is wired the same way on every platform"
+    )
+
+    methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
+    for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items():
+        method = methods.get(method_name)
+        assert method is not None, f"{path}::{class_name}.{method_name} is missing; it must call {required_call}"
+        calls = _method_calls(method)
+        assert required_call in calls, f"{path}::{class_name}.{method_name} must call self.{required_call}"
+
+
+@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
+def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: str):
+    """Ordering is load-bearing: the hook must run before the sampler reads the logits."""
+    import ast
+
+    classdef, path = _ar_runner_classdefs()[class_name]
+    sample = next(node for node in classdef.body if isinstance(node, ast.FunctionDef) and node.name == "_sample")
+
+    apply_lineno = _method_calls(sample).get("_apply_duplex_sampling")
+    sampler_lineno = next(
+        (
+            node.lineno
+            for node in ast.walk(sample)
+            if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == _MODEL_SAMPLER_CALL
+        ),
+        None,
+    )
+    assert apply_lineno is not None, f"{path}::{class_name}._sample must call self._apply_duplex_sampling"
+    assert sampler_lineno is not None, (
+        f"{path}::{class_name}._sample no longer calls {_MODEL_SAMPLER_CALL}(); update _MODEL_SAMPLER_CALL "
+        f"so the ordering check keeps testing something"
+    )
+    assert apply_lineno < sampler_lineno, (
+        f"{path}::{class_name}._sample calls self._apply_duplex_sampling at line {apply_lineno}, after "
+        f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}.  The hook masks the logits and publishes the "
+        f"row -> session map the sampler reads, so running it afterwards is a silent no-op."
+    )

vllm_omni/experimental/fullduplex/model_executor.py

@@ -4,7 +4,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import torch
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.sample.metadata import SamplingMetadata
 
 
 @dataclass(frozen=True, slots=True)
@@ -103,4 +108,69 @@ class DuplexSamplingHelper:
         self.hook_active = False
 
 
-__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow"]
+class DuplexSamplingRunnerMixin:
+    """Runner-side wiring for the experimental duplex sampling hook.
+
+    The GPU and NPU AR runners are siblings, not parent and child: ``GPUARModelRunner``
+    derives from ``OmniGPUModelRunner`` while ``NPUARModelRunner`` derives from
+    ``OmniNPUModelRunner``, so neither inherits the other's ``_sample``.  Keeping the
+    wiring here means a runner opts in by adding this mixin and calling the four
+    methods below, instead of re-deriving the hook logic and silently drifting.
+
+    The drift this exists to prevent already happened once: the NPU ``_sample`` was a
+    line-for-line copy of the GPU one minus the ``prepare_duplex_sampling`` call, which
+    left ``_minicpmo45_duplex_row_sessions`` empty on Ascend.  With no row -> session
+    map the MiniCPM-o 4.5 thinker could not tell that a turn was still in progress, so
+    every mid-turn ``<|listen|>`` was emitted verbatim instead of being rewritten to
+    ``<|tts_bos|>`` and the model never spoke again after the first turn.
+    """
+
+    def _init_duplex_sampling_state(self) -> None:
+        """Reset hook resolution.  Call from ``__init__``, before ``load_model``."""
+        self._duplex_sampling_hook = None
+        self._duplex_sampling_hook_resolved = False
+
+    def _resolve_duplex_sampling_hook(self, *, force: bool = False):
+        """Bind ``model.prepare_duplex_sampling``.  Call with ``force`` after load."""
+        if not force and getattr(self, "_duplex_sampling_hook_resolved", False):
+            return self._duplex_sampling_hook
+        candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None)
+        self._duplex_sampling_hook = candidate if callable(candidate) else None
+        self._duplex_sampling_hook_resolved = True
+        if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"):
+            self._duplex_sampling_helper = DuplexSamplingHelper()
+        return self._duplex_sampling_hook
+
+    def _update_duplex_sampling_states(self, scheduler_output: SchedulerOutput) -> None:
+        """Refresh which batch rows are duplex rows.  Call from ``_update_states``."""
+        if self._resolve_duplex_sampling_hook() is None:
+            return
+        helper = getattr(self, "_duplex_sampling_helper", None)
+        if helper is not None:
+            helper.update_states(self, scheduler_output)
+
+    def _apply_duplex_sampling(self, logits: torch.Tensor, prepared_sampling_metadata: SamplingMetadata) -> None:
+        """Hand the model its duplex rows.  Call from ``_sample``, before the sampler.
+
+        ``hook_active`` makes the call fire once more after the last duplex row leaves
+        the batch, so the model can drop the stale row -> session mapping without the
+        helper having to scan request state on every non-duplex step.
+        """
+        prepare_duplex_sampling = self._resolve_duplex_sampling_hook()
+        if prepare_duplex_sampling is None:
+            return
+        helper = getattr(self, "_duplex_sampling_helper", None)
+        rows = helper.rows(self) if helper is not None and helper.active_request_ids else ()
+        if rows or (helper is not None and helper.hook_active):
+            prepare_duplex_sampling(logits, prepared_sampling_metadata, rows)
+        if helper is not None:
+            helper.hook_active = bool(rows)
+
+    def _clear_duplex_sampling(self) -> None:
+        """Drop tracked rows.  Call from teardown paths that release device memory."""
+        helper = getattr(self, "_duplex_sampling_helper", None)
+        if helper is not None:
+            helper.clear()
+
+
+__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow", "DuplexSamplingRunnerMixin"]

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -44,6 +44,7 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture
 from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role, stage_sends_async_output
+from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
@@ -102,7 +103,7 @@ class ExecuteModelState(NamedTuple):
     batch_desc: BatchDescriptor
     multimodal_outputs: Any # Omni-Specific
 
-class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
+class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin):
     """Autoregressive NPU model runner that returns hidden states per request."""
 
     def __init__(self, *args, **kwargs):
@@ -140,6 +141,16 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
                 kv_transfer_manager=self.kv_transfer_manager,
             )
         self._downstream_payload_cache: dict[str, bool] = {}
+        self._init_duplex_sampling_state()
+
+    def load_model(self, *args, **kwargs) -> None:
+        super().load_model(*args, **kwargs)
+        self._resolve_duplex_sampling_hook(force=True)
+
+    def _update_states(self, scheduler_output: SchedulerOutput):
+        deferred_state_corrections_fn = super()._update_states(scheduler_output)
+        self._update_duplex_sampling_states(scheduler_output)
+        return deferred_state_corrections_fn
 
     def _make_buffer(self, *size, dtype, numpy=True):
         # Prevent ray from pinning the buffer due to large size
@@ -875,10 +886,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
                         self.input_batch.idx_mapping_np,
                         self.input_batch.positions[self.input_batch.logits_indices],
                     )
-                sampler_output = model_sample(
-                    logits,
-                    self._sampling_metadata_for_model_sampler(sampling_metadata),
-                )
+                prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
+                self._apply_duplex_sampling(logits, prepared_sampling_metadata)
+                sampler_output = model_sample(logits, prepared_sampling_metadata)
                 if sampler_output is not None:
                     return sampler_output
             return self.sampler(

vllm_omni/worker/gpu_ar_model_runner.py

@@ -46,6 +46,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import (
     get_stage_connector_role,
     stage_sends_async_output,
 )
+from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
@@ -290,7 +291,7 @@ class ExecuteModelState(NamedTuple):
     slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None
 
 
-class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
+class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin):
     """Autoregressive GPU model runner that returns hidden states per request.
 
     Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and
@@ -334,25 +335,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 kv_transfer_manager=self.kv_transfer_manager,
             )
         self._downstream_payload_cache: dict[str, bool] = {}
-        self._duplex_sampling_hook = None
-        self._duplex_sampling_hook_resolved = False
+        self._init_duplex_sampling_state()
 
     def load_model(self, *args, **kwargs) -> None:
         super().load_model(*args, **kwargs)
         self._resolve_duplex_sampling_hook(force=True)
 
-    def _resolve_duplex_sampling_hook(self, *, force: bool = False):
-        if not force and getattr(self, "_duplex_sampling_hook_resolved", False):
-            return self._duplex_sampling_hook
-        candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None)
-        self._duplex_sampling_hook = candidate if callable(candidate) else None
-        self._duplex_sampling_hook_resolved = True
-        if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"):
-            from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingHelper
-
-            self._duplex_sampling_helper = DuplexSamplingHelper()
-        return self._duplex_sampling_hook
-
     def _make_buffer(self, *size, dtype, numpy=True):
         # Prevent ray from pinning the buffer due to large size
         from vllm_omni.distributed.ray_utils.utils import (
@@ -421,11 +409,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
     def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None:
         deferred_state_corrections_fn = super()._update_states(scheduler_output)
-        if self._resolve_duplex_sampling_hook() is None:
-            return deferred_state_corrections_fn
-        helper = getattr(self, "_duplex_sampling_helper", None)
-        if helper is not None:
-            helper.update_states(self, scheduler_output)
+        self._update_duplex_sampling_states(scheduler_output)
         return deferred_state_corrections_fn
 
     def _request_final_stage_id(self, req_id: str) -> int | None:
@@ -553,9 +537,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             self._downstream_payload_cache.clear()
         if hasattr(self, "model_intermediate_buffer"):
             self.model_intermediate_buffer.clear()
-        duplex_helper = getattr(self, "_duplex_sampling_helper", None)
-        if duplex_helper is not None:
-            duplex_helper.clear()
+        self._clear_duplex_sampling()
 
         # 5. Release all CUDA graphs unconditionally (upstream only does this
         #    on ROCm; on CUDA the graphs are only freed by Python GC during
@@ -1468,14 +1450,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         self.input_batch.positions[self.input_batch.logits_indices],
                     )
                 prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata)
-                prepare_duplex_sampling = self._resolve_duplex_sampling_hook()
-                if prepare_duplex_sampling is not None:
-                    helper = getattr(self, "_duplex_sampling_helper", None)
-                    rows = helper.rows(self) if helper is not None and helper.active_request_ids else ()
-                    if rows or (helper is not None and helper.hook_active):
-                        prepare_duplex_sampling(logits, prepared_sampling_metadata, rows)
-                    if helper is not None:
-                        helper.hook_active = bool(rows)
+                self._apply_duplex_sampling(logits, prepared_sampling_metadata)
                 sampler_output = model_sample(logits, prepared_sampling_metadata)
                 if sampler_output is not None:
                     return sampler_output