From 009b80d686febcf683fdbc2bcdf3ad752884641e Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Fri, 7 Aug 2026 22:38:45 +0800 Subject: [PATCH] [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 --- tests/worker/test_native_duplex_hooks.py | 162 ++++++++++++++++++ .../experimental/fullduplex/model_executor.py | 74 +++++++- .../npu/worker/npu_ar_model_runner.py | 20 ++- vllm_omni/worker/gpu_ar_model_runner.py | 37 +--- 4 files changed, 255 insertions(+), 38 deletions(-) diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 383379202..42e691cca 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/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." + ) diff --git a/vllm_omni/experimental/fullduplex/model_executor.py b/vllm_omni/experimental/fullduplex/model_executor.py index eebe01876..540df8093 100644 --- a/vllm_omni/experimental/fullduplex/model_executor.py +++ b/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"] diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 1162eb088..c5d9468fe 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index ca7c2a21a..f1384968d 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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