vllm-omni/vllm_omni/diffusion/diffusion_kv/kv_connector.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Assemble native vLLM KV connectors for Scheduler-owned diffusion pages."""

from __future__ import annotations

import math
import time
from collections import Counter
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any

from vllm.config import KVTransferConfig, VllmConfig
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
from vllm.distributed.kv_transfer.kv_transfer_state import ensure_kv_transfer_initialized, ensure_kv_transfer_shutdown
from vllm.logger import init_logger
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.outputs import KVConnectorOutput

if TYPE_CHECKING:
    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
    from vllm.v1.core.kv_cache_manager import KVCacheManager
    from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector

    from vllm_omni.diffusion.data import OmniDiffusionConfig
    from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput

logger = init_logger(__name__)


class KVTransferRegistrationError(ValueError):
    """A failed handoff whose destination pages have not been dispatched."""


def mint_transfer_id(request_id: str) -> str:
    """Return the stable ticket shared by one AR -> Diffusion handoff."""

    return f"xfer-{request_id}"


def build_source_kv_transfer_params(
    *,
    transfer_id: str,
    remote_engine_id: str | None,
    remote_bootstrap_addr: str | None,
) -> dict[str, Any]:
    """Build the opaque producer-side metadata bag for an AR request."""

    params: dict[str, Any] = {
        "transfer_id": transfer_id,
        "do_remote_decode": True,
        "do_remote_prefill": False,
    }
    if remote_engine_id:
        params["remote_engine_id"] = remote_engine_id
    if remote_bootstrap_addr:
        params["remote_bootstrap_addr"] = remote_bootstrap_addr
    return params


def build_target_kv_transfer_params(
    *,
    source_params: Mapping[str, Any],
    remote_engine_id: str | None,
    remote_bootstrap_addr: str | None,
) -> dict[str, Any]:
    """Build the consumer-side metadata bag without interpreting pages."""

    params = dict(source_params)
    params["do_remote_prefill"] = True
    params["do_remote_decode"] = False
    if remote_engine_id:
        params["remote_engine_id"] = remote_engine_id
    if remote_bootstrap_addr:
        params["remote_bootstrap_addr"] = remote_bootstrap_addr
    return params


def bootstrap_addr_from_kv_transfer_config(kv_transfer_config: KVTransferConfig | None) -> str | None:
    """Read an optional connector bootstrap endpoint from native config."""

    if kv_transfer_config is None:
        return None
    extra_config = kv_transfer_config.kv_connector_extra_config or {}
    for key in ("bootstrap_addr", "prefill_bootstrap_addr", "mooncake_master"):
        value = extra_config.get(key)
        if isinstance(value, str) and value.strip():
            return value
    return None


def parse_kv_transfer_config(value: object | None) -> KVTransferConfig | None:
    """Normalize a YAML mapping or an existing native config."""

    if value is None:
        return None
    if isinstance(value, KVTransferConfig):
        _validate_kv_transfer_config(value)
        return value
    if isinstance(value, Mapping):
        payload = dict(value)
        if not payload:
            return None
        engine_id = payload.get("engine_id")
        if not isinstance(engine_id, str) or not engine_id.strip():
            raise ValueError("Diffusion native kv_transfer_config requires a non-empty engine_id")
        config = KVTransferConfig(**payload)
        _validate_kv_transfer_config(config)
        return config
    raise TypeError(f"kv_transfer_config must be a mapping or KVTransferConfig, got {type(value)!r}")


def create_scheduler_kv_connector(
    od_config: OmniDiffusionConfig,
    kv_cache_config: KVCacheConfig | None = None,
    vllm_config: VllmConfig | None = None,
) -> KVConnectorBase_V1 | None:
    """Create a Scheduler-role connector when native config is present."""

    kv_transfer_config = getattr(od_config, "kv_transfer_config", None)
    if kv_transfer_config is None:
        return None
    if not isinstance(kv_transfer_config, KVTransferConfig):
        raise TypeError(
            f"Diffusion native kv_transfer_config must be KVTransferConfig, got {type(kv_transfer_config)!r}"
        )

    if kv_cache_config is None or vllm_config is None:
        raise ValueError("KV transfer requires paged_scheduler and a native cache configuration")
    connector = KVConnectorFactory.create_connector(
        config=vllm_config,
        role=KVConnectorRole.SCHEDULER,
        kv_cache_config=kv_cache_config,
    )
    logger.info(
        "Created KV connector (SCHEDULER role): connector=%s engine_id=%s",
        kv_transfer_config.kv_connector,
        kv_transfer_config.engine_id,
    )
    return connector


def init_worker_kv_connector(vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
    """Initialize the Worker-role connector with its rank-local cache plan."""

    if vllm_config.kv_transfer_config is None:
        return
    import vllm.distributed.parallel_state as parallel_state

    from vllm_omni.diffusion.distributed.parallel_state import get_sp_group

    # vLLM 0.29 Mooncake snapshots tp_rank/tp_size in its Worker constructor;
    # consumer receive/region mapping use those saved values, not get_tp_group.
    # Restrict this compatibility bridge to that connector and restore the
    # model's TP group even when construction fails. Never remap model TP for
    # the connector's entire lifetime (it would change model collectives).
    tp_group = parallel_state._TP
    parallel_config = getattr(vllm_config, "parallel_config", None)
    if getattr(parallel_config, "prefill_context_parallel_size", 1) > 1 and tp_group.world_size == 1:
        if vllm_config.kv_transfer_config.kv_connector != "MooncakeConnector":
            raise ValueError("Native SP KV transfer currently requires MooncakeConnector")
        parallel_state._TP = get_sp_group()
    try:
        ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
    finally:
        parallel_state._TP = tp_group
    logger.info(
        "Initialized KV connector (WORKER role): connector=%s engine_id=%s",
        vllm_config.kv_transfer_config.kv_connector,
        vllm_config.kv_transfer_config.engine_id,
    )


def shutdown_kv_connector(*, scheduler_connector: KVConnectorBase_V1 | None = None) -> None:
    """Shutdown Worker and Scheduler connector objects idempotently."""

    if scheduler_connector is not None:
        scheduler_connector.shutdown()
    ensure_kv_transfer_shutdown()


def prepare_kv_requests(requests: tuple[DiffusionKVRequest, ...], params: Mapping[str, Any]) -> None:
    for request in requests:
        request.kv_transfer_params = dict(params)
        num_tokens = min(len(request.prompt_token_ids), params["num_transfer_tokens"])
        request.prompt_token_ids = request.prompt_token_ids[:num_tokens]
        request.num_prompt_tokens = num_tokens


def install_mooncake_cfg_fanout(connector: KVConnectorBase_V1) -> None:
    """Account for multiple CFG destinations on one native Mooncake ticket.

    vLLM 0.29 initializes ``need_send`` to the paired consumer rank count,
    but increments ``sent`` per consumer request. Adapt only this producer
    instance, after KV initialization and before serving requests. Keep the
    upstream address planning, transport, timeout and completion code intact.
    """
    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import MooncakeConnector

    if not isinstance(connector, MooncakeConnector):
        return
    worker = connector.connector_worker
    if worker is None or worker.is_kv_consumer or getattr(worker, "_omni_cfg_fanout_installed", False):
        return
    build_transfer_params = worker._build_transfer_params

    async def build_with_fanout(ready_reqs, agent_meta, local_regions, remote_regions):
        # commit_kv_load registers all rows before build_connector_meta. Thus
        # each rank's message contains the complete row set for each ticket,
        # including empty-block notifications. Count the full message, not
        # just the subset whose producer-ready events fired in this iteration.
        rows_per_ticket = Counter(transfer_id for transfer_id, _ in agent_meta.req_blocks.values())
        tickets = {send_meta.transfer_id: send_meta for _, send_meta in ready_reqs}
        for transfer_id, send_meta in tickets.items():
            rows = rows_per_ticket[transfer_id]
            previous = getattr(send_meta, "_omni_cfg_rows", None)
            if rows < 1 or (previous is not None and previous != rows):
                # Use the upstream transfer-error path: it decrements sending
                # without incrementing sent, so a malformed rank cannot free
                # source pages still needed by another rank.
                return [], [], [], [req_id for req_id, _ in ready_reqs], "Inconsistent Mooncake CFG receive count"
        # No await between checking and updating: all handlers run on the
        # sender loop. The marker belongs to the ticket, not to a request-global
        # variable; it disappears with the upstream ticket on completion/expiry.
        for transfer_id, send_meta in tickets.items():
            if getattr(send_meta, "_omni_cfg_rows", None) is None:
                send_meta.need_send *= rows_per_ticket[transfer_id]
                send_meta._omni_cfg_rows = rows_per_ticket[transfer_id]
        return await build_transfer_params(ready_reqs, agent_meta, local_regions, remote_regions)

    worker._build_transfer_params = build_with_fanout  # type: ignore[method-assign]
    setattr(worker, "_omni_cfg_fanout_installed", True)


def commit_kv_load(
    connector: KVConnectorBase_V1,
    manager: KVCacheManager,
    requests: tuple[DiffusionKVRequest, ...],
    matched_tokens: list[int],
) -> set[str]:
    expected = set()
    allocations = []
    # All rows sharing a transfer_id must reach the same connector metadata:
    # the producer uses this complete per-rank row set for fan-out accounting.
    # Validate every CFG row before mutating any connector state.
    for request, num_tokens in zip(requests, matched_tokens, strict=True):
        blocks = manager.get_blocks(request.request_id)
        # Mooncake's producer advertises complete physical blocks. Keep every
        # CFG receiver on that same block boundary so the native connector can
        # use its normal one-to-one page mapping even when the branches have
        # different logical prefix lengths. Bytes beyond ``num_tokens`` are
        # not marked computed and are overwritten by the DiT prefill.
        transfer_tokens = num_tokens
        if num_tokens > 0 and request.kv_transfer_params is not None:
            transfer_tokens = request.kv_transfer_params["num_transfer_tokens"]
            if type(transfer_tokens) is not int or not num_tokens <= transfer_tokens <= request.num_tokens:
                raise KVTransferRegistrationError(
                    "Diffusion KV transfer boundary must cover the reusable prefix "
                    f"without exceeding the allocated sequence: reusable={num_tokens}, "
                    f"transfer={transfer_tokens!r}, allocated={request.num_tokens}"
                )
        prefix_blocks = KVCacheBlocks(
            tuple(
                group[: (transfer_tokens + spec.kv_cache_spec.block_size - 1) // spec.kv_cache_spec.block_size]
                for group, spec in zip(blocks.blocks, manager.kv_cache_config.kv_cache_groups, strict=True)
            )
        )
        allocations.append((request, prefix_blocks, num_tokens))
    try:
        for request, prefix_blocks, num_tokens in allocations:
            connector.update_state_after_alloc(request, prefix_blocks, num_tokens)
            request.num_computed_tokens = num_tokens
            if request.kv_transfer_params is not None:
                expected.add(request.request_id)
    except Exception as exc:
        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import MooncakeConnector
        from vllm.v1.request import RequestStatus

        if not isinstance(connector, MooncakeConnector):
            # Unknown connectors may have started I/O during registration.
            raise
        # Mooncake only queues metadata here; no Worker has seen addresses yet.
        # Re-arm its pre-scheduling abort hook to replace even partially
        # registered CFG receives with empty-block notifications to the producer.
        for request in requests:
            if request.kv_transfer_params is not None:
                request.kv_transfer_params["do_remote_prefill"] = True
                request.status = RequestStatus.FINISHED_ABORTED
                delay_free, _ = connector.request_finished(request, [])
                if delay_free:
                    raise RuntimeError("Connector could not cancel an undispatched KV load") from exc
        raise KVTransferRegistrationError(f"Could not register diffusion KV receive: {exc}") from exc
    return expected


def wait_for_kv_load(
    active_connector: ActiveKVConnector, scheduler_output: DiffusionSchedulerOutput, timeout: float
) -> KVConnectorOutput:
    connector = active_connector.kv_connector
    active_connector.pre_forward(scheduler_output)
    pending = set(scheduler_output.kv_transfer_request_ids)
    finished_ids = scheduler_output.kv_finished_request_ids
    if not pending:
        output = active_connector.post_forward(finished_ids)
        if output.invalid_block_ids:
            raise RuntimeError("Diffusion KV connector reported invalid remote pages")
        return output

    received, sent = set(), set()
    deadline = time.monotonic() + timeout
    while True:
        finished_sending, finished_recving = connector.get_finished(finished_ids)
        sent.update(finished_sending or ())
        received.update(finished_recving or ())
        pending.difference_update(received)
        if connector.get_block_ids_with_load_errors():
            raise RuntimeError("Diffusion KV connector reported invalid remote pages")
        if not pending:
            break
        if time.monotonic() >= deadline:
            # A timeout is not a cancellation acknowledgement. Return partial
            # completion; Scheduler quarantines the remaining destination pages.
            break
        time.sleep(0.001)
    output = active_connector.post_forward(finished_ids)
    output.finished_sending = sent | (output.finished_sending or set())
    output.finished_recving = received | (output.finished_recving or set())
    if output.invalid_block_ids:
        raise RuntimeError("Diffusion KV connector reported invalid remote pages")
    return output


def _validate_kv_transfer_config(config: KVTransferConfig) -> None:
    engine_id = config.engine_id
    if not isinstance(engine_id, str) or not engine_id.strip():
        raise ValueError("Diffusion native kv_transfer_config requires a non-empty engine_id")
    extra_config = config.kv_connector_extra_config or {}
    timeout = extra_config.get("transfer_timeout", 60.0)
    if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0:
        raise ValueError("Diffusion transfer_timeout must be a positive finite number of seconds")
    if config.kv_connector is not None and config.kv_role is None:
        raise ValueError("Diffusion native kv_transfer_config requires kv_role when kv_connector is set")
    if config.kv_role not in (None, "kv_consumer", "kv_producer", "kv_both"):
        raise ValueError(f"Unsupported KV connector role: {config.kv_role!r}")