vllm-omni/vllm_omni/model_executor/duplex_sampling.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

from dataclasses import dataclass
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)
class DuplexSamplingRow:
    """Request-local context for the experimental duplex sampling hook."""

    row_idx: int
    request_id: str
    session_id: str | None
    seq: int | None
    payload: dict[str, object] | None
    max_tokens: int | None


class DuplexSamplingHelper:
    """Tracks duplex sampling rows outside the generic AR model runner."""

    def __init__(self) -> None:
        self.active_request_ids: set[str] = set()
        self.hook_active = False

    @staticmethod
    def _is_duplex_data_plane_info(info: object) -> bool:
        duplex = info.get("duplex") if isinstance(info, dict) else None
        return isinstance(duplex, dict) and duplex.get("data_plane") is True

    @staticmethod
    def _request_intermediate_info(runner: object, req_id: str) -> dict[str, Any] | None:
        model_intermediate_buffer = getattr(runner, "model_intermediate_buffer", {})
        info = model_intermediate_buffer.get(req_id) if isinstance(model_intermediate_buffer, dict) else None
        if isinstance(info, dict):
            return info
        requests = getattr(runner, "requests", {})
        req_state = requests.get(req_id) if isinstance(requests, dict) else None
        info = getattr(req_state, "additional_information_cpu", None)
        return info if isinstance(info, dict) else None

    def refresh_active_request(self, runner: object, req_id: str) -> None:
        if self._is_duplex_data_plane_info(self._request_intermediate_info(runner, req_id)):
            self.active_request_ids.add(req_id)
        else:
            self.active_request_ids.discard(req_id)

    def update_states(self, runner: object, scheduler_output: object) -> None:
        self.active_request_ids.difference_update(str(req_id) for req_id in scheduler_output.finished_req_ids)
        for request in scheduler_output.scheduled_new_reqs:
            self.refresh_active_request(runner, str(request.req_id))

    def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]:
        rows: list[DuplexSamplingRow] = []
        req_ids = [str(req_id) for req_id in getattr(runner.input_batch, "req_ids", [])]
        requests = getattr(runner, "requests", {})
        for row_idx, req_id in enumerate(req_ids):
            if req_id not in self.active_request_ids:
                continue
            info = self._request_intermediate_info(runner, req_id)
            duplex = info.get("duplex") if isinstance(info, dict) else None
            if not isinstance(duplex, dict) or duplex.get("data_plane") is not True:
                continue
            session_id = duplex.get("session_id")
            if not isinstance(session_id, str) or not session_id:
                session_id = None
            try:
                seq = int(duplex.get("seq"))
            except (TypeError, ValueError):
                seq = None
            payload = duplex.get("payload")
            if not isinstance(payload, dict):
                payload = None
            request = requests.get(req_id) if isinstance(requests, dict) else None
            sampling_params = getattr(request, "sampling_params", None)
            try:
                max_tokens = int(getattr(sampling_params, "max_tokens", 0) or 0)
            except (TypeError, ValueError):
                max_tokens = 0
            rows.append(
                DuplexSamplingRow(
                    row_idx=row_idx,
                    request_id=req_id,
                    session_id=session_id,
                    seq=seq,
                    payload=payload,
                    max_tokens=max_tokens if max_tokens > 0 else None,
                )
            )
        return tuple(rows)

    def clear(self) -> None:
        self.active_request_ids.clear()
        self.hook_active = False


class DuplexSamplingRunnerMixin:
    """Runner-side wiring for the experimental duplex sampling hook.

    GPU and NPU AR runners are siblings, not parent and child, so neither
    inherits the other's ``_sample``. Keep the hook here and have each runner
    opt in, instead of copying the GPU ``_sample`` path and dropping the
    ``prepare_duplex_sampling`` call.
    """

    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."""
        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:
        helper = getattr(self, "_duplex_sampling_helper", None)
        if helper is not None:
            helper.clear()


__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow", "DuplexSamplingRunnerMixin"]