[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>

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