vllm-omni/vllm_omni/experimental/ar_diffusion/runner.py

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

import dataclasses
import time
from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager

import torch
from vllm.logger import init_logger

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.models.interface import supports_step_execution
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput, KVPrefetchJob
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
from vllm_omni.experimental.ar_diffusion.capability import (
    ARDiffusionKVCacheSpec,
    SupportsARDiffusionPipeline,
    SupportsARDiffusionWarmup,
    supports_chunk_step_grouping,
)
from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
from vllm_omni.platforms import current_omni_platform

logger = init_logger(__name__)


def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
    """Resolve optional deployment overrides and enable AR-Diffusion KV."""
    raw = getattr(od_config, "ar_diffusion_kv_config", None)
    if raw is None:
        model_config = getattr(od_config, "model_config", None)
        if isinstance(model_config, dict):
            raw = model_config.get("ar_diffusion_kv_config")
    if isinstance(raw, ARDiffusionKVConfig):
        return dataclasses.replace(raw, enable=True)
    if isinstance(raw, dict):
        return ARDiffusionKVConfig(**{**raw, "enable": True})
    return ARDiffusionKVConfig(enable=True)


class ARDiffusionModelRunner(DiffusionModelRunner):
    """Own AR KV pools and session lifetime for a capable diffusion pipeline.

    The runner never inspects model internals. At load time it consumes an
    :class:`ARDiffusionKVCacheSpec`; during execution it binds a runner-owned
    :class:`ARDiffusionKVState` through the pipeline's context manager. Sessions
    persist across requests and end on explicit close, LRU eviction, or failed
    forward. A request with ``extra_args["reset"]`` releases the old session
    before executing with a fresh one.
    """

    _WARMUP_SID = "__ardiffusion_warmup__"

    def __init__(self, vllm_config: object, od_config: OmniDiffusionConfig, device: torch.device) -> None:
        super().__init__(vllm_config, od_config, device)
        self.ar_diffusion_kv_config = resolve_ar_diffusion_kv_config(od_config)
        self.kv_cache: ARDiffusionKVCache | None = None
        self._ar_diffusion_capability: SupportsARDiffusionPipeline | None = None
        self._ar_diffusion_kv_cache_spec: ARDiffusionKVCacheSpec | None = None
        self._sessions: OrderedDict[str, ARDiffusionKVState] = OrderedDict()
        self._session_capacity = 0
        self._perf_e2e_times: list[float] = []
        # Wall-clock start of each request's in-flight stepwise chunk, so
        # ``_perf_e2e_times`` keeps one entry per AR block in both execution
        # modes instead of one entry per denoise step.
        self._stepwise_chunk_started: dict[str, float] = {}

    @staticmethod
    def _require_capability(pipeline: object) -> SupportsARDiffusionPipeline:
        if not isinstance(pipeline, SupportsARDiffusionPipeline):
            raise TypeError(
                "ARDiffusionEngine requires the loaded pipeline to implement "
                "SupportsARDiffusionPipeline: ar_diffusion_kv_cache_spec(), "
                "bind_ar_diffusion_state(), reset_ar_diffusion_session(), and "
                "close_ar_diffusion_session(). Select the default DiffusionEngine "
                "or add an explicit AR-Diffusion capability adapter to the model."
            )
        return pipeline

    def load_model(self, *args: object, **kwargs: object) -> None:
        super().load_model(*args, **kwargs)
        if self.pipeline is None:
            return
        self._preallocate_kv_cache()
        if not self.od_config.enforce_eager and self.ar_diffusion_kv_config.warmup_cudagraph:
            self._warmup_ar_rollout()

    def _available_memory_bytes(self) -> int:
        if self.device is None:
            raise RuntimeError("AR-Diffusion KV preallocation requires an initialized device")
        return current_omni_platform.get_free_memory(torch.device(self.device))

    def _preallocate_kv_cache(self, *, available_bytes: int | None = None) -> None:
        """Build pools solely from the pipeline capability and runner config."""
        if bool(getattr(self.od_config, "step_execution", False)) and not supports_step_execution(self.pipeline):
            raise ValueError(
                "ARDiffusionModelRunner currently supports request-mode execution only; "
                "step_execution=True would bypass per-request AR session binding."
            )
        max_num_seqs = int(getattr(self.od_config, "max_num_seqs", 1) or 1)
        if max_num_seqs > 1:
            raise ValueError(
                "AR-Diffusion paged KV supports max_num_seqs=1 (single-sequence "
                f"rollouts); got max_num_seqs={max_num_seqs}."
            )
        capability = self._require_capability(self.pipeline)
        if bool(getattr(self.pipeline, "supports_request_batch", False)):
            raise ValueError(
                "ARDiffusionModelRunner currently supports one request at a time; "
                "request-batch execution would bypass per-request AR session binding."
            )
        spec = capability.ar_diffusion_kv_cache_spec()
        if not isinstance(spec, ARDiffusionKVCacheSpec):
            raise TypeError(
                f"ar_diffusion_kv_cache_spec() must return ARDiffusionKVCacheSpec, got {type(spec).__name__}"
            )

        config = dataclasses.replace(
            self.ar_diffusion_kv_config,
            chunk_size=spec.tokens_per_frame,
            window_chunks=self.ar_diffusion_kv_config.window_chunks or spec.window_frames,
            sink_chunks=self.ar_diffusion_kv_config.sink_chunks or spec.sink_frames,
            reset_at_boundary=self.ar_diffusion_kv_config.reset_at_boundary or spec.reset_at_boundary,
        )
        self.ar_diffusion_kv_config = config
        self._ar_diffusion_capability = capability
        self._ar_diffusion_kv_cache_spec = spec
        self.kv_cache = ARDiffusionKVCache(
            config,
            num_layers=spec.num_layers,
            num_kv_heads=spec.num_kv_heads,
            head_size=spec.head_size,
            dtype=self.od_config.dtype,
            block_size=spec.tokens_per_frame,
            max_model_len=spec.max_model_len,
            available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes,
            kv_branches=spec.kv_branches,
            session_capacity=spec.session_capacity,
            cross_attention_lengths=spec.cross_attention_lengths,
            cross_attention_kv_heads=spec.cross_attention_kv_heads,
            frames_per_block=spec.frames_per_block,
            max_scratch_tokens_per_branch=spec.max_scratch_tokens_per_branch,
            model_owned_state_bytes_per_session=spec.model_owned_state_bytes_per_session,
            device=self.device,
        )
        self._session_capacity = self.kv_cache.session_capacity
        logger.info(
            "AR-Diffusion KV cache: blocks=%d layers=%d local_kv_heads=%d head_size=%d "
            "tokens/frame=%d frames/block=%d window=%d sink=%d kv_branches=%s cross=%s "
            "resident_capacity=%d requested_capacity=%d",
            self.kv_cache.num_blocks,
            spec.num_layers,
            spec.num_kv_heads,
            spec.head_size,
            spec.tokens_per_frame,
            spec.frames_per_block,
            config.window_chunks,
            config.sink_chunks,
            [(kv_branch.name, kv_branch.local_index) for kv_branch in spec.kv_branches],
            {
                name: (length, spec.cross_attention_kv_heads.get(name, spec.num_kv_heads))
                for name, length in spec.cross_attention_lengths.items()
            },
            self._session_capacity,
            spec.session_capacity,
        )

    def _new_session_state(self, session_id: str) -> ARDiffusionKVState:
        if self.kv_cache is None or self._ar_diffusion_kv_cache_spec is None:
            raise RuntimeError("AR-Diffusion session requested before KV cache initialization")
        adapters = {
            kv_branch.name: self.kv_cache.begin_request(f"ar::{session_id}::{kv_branch.name}")
            for kv_branch in self._ar_diffusion_kv_cache_spec.kv_branches
        }
        return ARDiffusionKVState(
            self.kv_cache,
            session_id,
            adapters,
            num_layers=self._ar_diffusion_kv_cache_spec.num_layers,
        )

    def _release_session(
        self,
        session_id: str,
        *,
        reset_model: bool,
        reason: str,
        suppress_errors: bool = False,
    ) -> None:
        """One release path for reset, close, eviction, and failed forwards."""
        state = self._sessions.pop(session_id, None)
        errors: list[Exception] = []
        if state is not None:
            try:
                state.close()
            except Exception as exc:  # noqa: BLE001 - notify the pipeline even if pool cleanup fails
                errors.append(exc)
        capability = self._ar_diffusion_capability
        if capability is not None:
            try:
                if reset_model:
                    capability.reset_ar_diffusion_session(session_id)
                else:
                    capability.close_ar_diffusion_session(session_id)
            except Exception as exc:  # noqa: BLE001 - preserve all lifecycle cleanup attempts
                errors.append(exc)
        logger.debug("AR-Diffusion released session=%s reason=%s", session_id, reason)
        if errors:
            if suppress_errors:
                logger.warning(
                    "AR-Diffusion session=%s cleanup after %s had %d error(s): %s",
                    session_id,
                    reason,
                    len(errors),
                    errors,
                )
            else:
                raise errors[0]

    def reset_session(self, session_id: str) -> None:
        """Release KV and notify the pipeline to reset model-owned state."""
        self._release_session(session_id, reset_model=True, reason="reset")

    def close_session(self, session_id: str) -> None:
        """Release KV and notify the pipeline to drop model-owned state."""
        self._release_session(session_id, reset_model=False, reason="close")

    def _get_or_create_session(self, session_id: str) -> ARDiffusionKVState:
        state = self._sessions.get(session_id)
        if state is None:
            while len(self._sessions) >= self._session_capacity:
                oldest = next(iter(self._sessions))
                self._release_session(oldest, reset_model=False, reason="lru_eviction")
            state = self._new_session_state(session_id)
            self._sessions[session_id] = state
        self._sessions.move_to_end(session_id)
        return state

    @staticmethod
    def _request_session(
        req: OmniDiffusionRequest,
    ) -> tuple[str, dict, ARDiffusionTickRequest | None]:
        extra_args = req.sampling_params.extra_args or {}
        tick = ARDiffusionTickRequest.from_extra_args(extra_args)
        if tick is not None:
            return tick.session_id, extra_args, tick
        return str(extra_args.get("session_id") or "default"), extra_args, None

    @contextmanager
    def _bound_ar_session(
        self,
        session_id: str,
        *,
        description: str,
        synchronize: bool = True,
    ) -> Iterator[None]:
        """Bind runner-owned KV for one runner invocation and fail closed.

        The pipeline sees the state only inside this context. On any failure the
        session is released before the exception propagates, so a later request
        for the same id cannot resume half-written pages.
        """
        capability = self._ar_diffusion_capability
        if capability is None:
            raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")
        state = self._get_or_create_session(session_id)
        try:
            with capability.bind_ar_diffusion_state(session_id, state):
                yield
            if synchronize:
                current_omni_platform.synchronize()
        except Exception:
            self._release_session(
                session_id,
                reset_model=False,
                reason="forward_exception",
                suppress_errors=True,
            )
            logger.warning(
                "AR-Diffusion %s failed for session=%s; KV and model state were released",
                description,
                session_id,
            )
            raise

    def _release_scheduler_finished_sessions(self, scheduler_output: DiffusionSchedulerOutput) -> None:
        """Close sessions the scheduler retired, including aborted requests.

        A stepwise request that runs to completion is closed as soon as its
        final chunk is emitted; an aborted one never produces that output, so
        its KV would otherwise linger until LRU eviction. Session ids equal
        request ids on this path, and tick sessions are keyed by their own ids,
        so the membership check leaves them untouched.
        """
        for request_id in getattr(scheduler_output, "finished_req_ids", None) or ():
            self._stepwise_chunk_started.pop(request_id, None)
            if request_id in self._sessions:
                self._release_session(
                    request_id,
                    reset_model=False,
                    reason="close",
                    suppress_errors=True,
                )

    def execute_model(
        self,
        req: OmniDiffusionRequest,
        kv_prefetch_job: KVPrefetchJob | None = None,
    ) -> DiffusionOutput:
        if self.kv_cache is None:
            return super().execute_model(req, kv_prefetch_job=kv_prefetch_job)
        if self._ar_diffusion_capability is None:
            raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")

        session_id, extra_args, tick = self._request_session(req)
        reset = tick.reset if tick is not None else bool(extra_args.get("reset", False))
        close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False))
        if reset:
            self.reset_session(session_id)
        started = time.perf_counter()
        with self._bound_ar_session(session_id, description="forward"):
            output = super().execute_model(req, kv_prefetch_job=kv_prefetch_job)
        self._perf_e2e_times.append(time.perf_counter() - started)
        if close_session:
            self.close_session(session_id)
        return output

    def execute_model_batch(
        self,
        scheduler_output: DiffusionSchedulerOutput,
        od_config: OmniDiffusionConfig,
    ) -> BatchRunnerOutput:
        """Reject request batching until batch-aware AR state binding exists."""
        raise RuntimeError(
            "ARDiffusionModelRunner does not support request-batch execution; use request mode with max_num_seqs=1."
        )

    def _groups_chunk_steps(self, scheduler_output: DiffusionSchedulerOutput) -> bool:
        """Policy: does this stepwise call run the scheduled request's chunk to its boundary?

        Only on the realtime AR path, for a lone streaming request of a
        pipeline that declares ``supports_chunk_step_grouping``. The session's
        KV is bound for the whole call and the scheduler regains control at
        every chunk boundary, which is where this path's interactions are
        applied anyway; what it gives up is the chance to act between two
        steps of one chunk.
        """
        return (
            bool(self.od_config.streaming_output)
            and len(scheduler_output.scheduled_request_ids) == 1
            and supports_chunk_step_grouping(self.pipeline)
        )

    def _execute_stepwise_core(
        self,
        scheduler_output: DiffusionSchedulerOutput,
        *,
        record_output_peak_memory: bool,
        in_diffusion_kv_memory_profile: bool = False,
    ) -> BatchRunnerOutput:
        """Run the shared single-step core once, or repeatedly until the chunk boundary.

        The shared runner keeps its one-step-per-call semantics. When the
        grouping policy applies, the remaining steps of the request's current
        chunk are driven from here with a continuation of the scheduler
        output: the request is presented as a cached request, with nothing
        new to admit, nothing to retire and no connector work, so every
        further step is exactly what the next scheduler cycle would have
        asked for. Grouping the probes and the commit of one block this way
        drops the scheduler/executor round trips between them.

        The loop stops at the first call that emits a result (a chunk or an
        error), finishes the request, or leaves it without state; the base
        core already handles a mid-chunk failure or interrupt by finishing
        the request with an error and dropping its state.
        """
        output = super()._execute_stepwise_core(
            scheduler_output,
            record_output_peak_memory=record_output_peak_memory,
            in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
        )
        if not self._groups_chunk_steps(scheduler_output):
            return output
        request_id = scheduler_output.scheduled_request_ids[0]
        continuation = dataclasses.replace(
            scheduler_output,
            scheduled_new_reqs=[],
            scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
            finished_req_ids=set(),
            kv_prefetch_job=None,
            kv_connector_metadata=None,
        )
        if not self._chunk_step_pending(output, request_id):
            return output
        max_steps = self.state_cache[request_id].chunk_num_steps
        if max_steps is None:
            return output
        steps = 1
        while self._chunk_step_pending(output, request_id):
            if steps >= max_steps:
                self.state_cache.pop(request_id)
                runner_output = output.get_request_output(request_id)
                assert runner_output is not None
                runner_output.finished = True
                runner_output.result = DiffusionOutput(error=f"Chunk did not complete within {max_steps} denoise steps")
                break
            steps += 1
            output = super()._execute_stepwise_core(
                continuation,
                record_output_peak_memory=record_output_peak_memory,
                in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
            )
        return output

    def _chunk_step_pending(self, output: BatchRunnerOutput, request_id: str) -> bool:
        """More steps of the request's current chunk remain after ``output``."""
        runner_output = output.get_request_output(request_id)
        if runner_output is None or runner_output.finished or runner_output.result is not None:
            return False
        return request_id in self.state_cache

    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
        """Bind runner-owned KV for one stepwise invocation, then inherit the step loop."""
        pipeline = getattr(self, "pipeline", None)
        if pipeline is None or not supports_step_execution(pipeline):
            raise RuntimeError(
                "ARDiffusionModelRunner does not support step execution; use request mode with step_execution=False."
            )
        if self.kv_cache is None:
            return super().execute_stepwise(scheduler_output)
        if self._ar_diffusion_capability is None:
            raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")
        # Retire sessions the scheduler dropped before serving this cycle: a
        # completed request is closed when its final chunk is emitted, an
        # aborted one only here.
        self._release_scheduler_finished_sessions(scheduler_output)
        request_ids = list(getattr(scheduler_output, "scheduled_request_ids", ()))
        if not request_ids:
            return super().execute_stepwise(scheduler_output)
        if len(request_ids) != 1:
            raise RuntimeError(
                "AR-Diffusion step execution supports one request at a time; "
                f"got {len(request_ids)} scheduled requests."
            )
        request_id = request_ids[0]
        session_id = request_id
        chunk_started = self._stepwise_chunk_started.setdefault(request_id, time.perf_counter())
        try:
            # Denoise steps inside one chunk need no device barrier; the chunk
            # boundary below is the meaningful one, and syncing per step would
            # cost one synchronize per denoise step instead of per AR block.
            with self._bound_ar_session(session_id, description="stepwise", synchronize=False):
                output = super().execute_stepwise(scheduler_output)
            runner_output = output.get_request_output(request_id)
            finished = bool(runner_output is not None and runner_output.finished)
            result = None if runner_output is None else runner_output.result
            errored = bool(result is not None and getattr(result, "error", None))
            chunk_emitted = result is not None or finished or errored
            if chunk_emitted:
                # The barrier is where an asynchronous device error surfaces, so
                # it must sit inside the same fail-closed scope as the step
                # itself rather than leave cleanup to the scheduler's next cycle.
                current_omni_platform.synchronize()
        except Exception:
            self._fail_closed_stepwise(request_id)
            raise
        # One entry per emitted AR block keeps ``_perf_e2e_times`` comparable
        # with request mode, where one execute_model() is one block.
        if chunk_emitted:
            self._perf_e2e_times.append(time.perf_counter() - chunk_started)
            self._stepwise_chunk_started.pop(request_id, None)
        if finished or errored:
            self._release_session(
                session_id,
                reset_model=False,
                reason="forward_exception" if errored else "close",
                suppress_errors=True,
            )
        return output

    def _fail_closed_stepwise(self, request_id: str) -> None:
        """Drop everything a failed stepwise invocation left behind.

        A failure inside the step is already released by ``_bound_ar_session``;
        a failure at the chunk barrier happens after that context closed, so the
        session is still present and is released here.
        """
        self._stepwise_chunk_started.pop(request_id, None)
        state_cache = getattr(self, "state_cache", None)
        if isinstance(state_cache, dict):
            state_cache.pop(request_id, None)
        if request_id in self._sessions:
            self._release_session(request_id, reset_model=False, reason="forward_exception", suppress_errors=True)
            logger.warning(
                "AR-Diffusion stepwise chunk barrier failed for session=%s; KV and model state were released",
                request_id,
            )

    def _warmup_ar_rollout(self) -> None:
        """Run model-provided warmup requests, or safely skip when absent."""
        if self.kv_cache is None or self.pipeline is None:
            return
        if not isinstance(self.pipeline, SupportsARDiffusionWarmup):
            logger.info("AR-Diffusion pipeline provides no warmup requests; skipping rollout warmup")
            return
        sid = self._WARMUP_SID
        free_before = self.kv_cache.manager.block_pool.get_num_free_blocks()
        try:
            for request in self.pipeline.ar_diffusion_warmup_requests(sid):
                self.execute_model(request)
        except Exception as exc:  # noqa: BLE001 - warmup must not make model load fail
            logger.warning("AR-Diffusion rollout warmup failed (%s); using lazy capture", exc)
        finally:
            self._release_session(
                sid,
                reset_model=False,
                reason="warmup_complete",
                suppress_errors=True,
            )
            self._perf_e2e_times.clear()
            free_after = self.kv_cache.manager.block_pool.get_num_free_blocks()
            if free_after != free_before:
                logger.warning(
                    "AR-Diffusion warmup did not restore the KV pool (free %d -> %d)",
                    free_before,
                    free_after,
                )