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

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

from __future__ import annotations

from typing import TYPE_CHECKING

import torch
from vllm.config import VllmConfig
from vllm.v1.core.kv_cache_utils import (
    generate_scheduler_kv_cache_config,
    get_kv_cache_configs,
    resolve_kv_cache_block_sizes,
)
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
from vllm_omni.platforms import current_omni_platform

if TYPE_CHECKING:
    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
    from vllm_omni.diffusion.request import OmniDiffusionRequest


def build_native_kv_cache_configs(
    vllm_config: VllmConfig,
    worker_specs: list[dict[str, KVCacheSpec]],
    available_memory: list[int],
    *,
    indexes_kv_by_block_stride: bool = False,
) -> tuple[list[KVCacheConfig], KVCacheConfig]:
    """Build rank-local and Scheduler-native configs using vLLM utilities.

    ``indexes_kv_by_block_stride`` carries the native backend requirement that
    0.29 removed from ``AttentionSpec``.  Every diffusion backend in tree
    declares ``False``, so the default reproduces current behaviour; a caller
    that learns otherwise from its workers passes ``True`` here and the
    block-outermost layout propagates to every rank-local config.
    """

    from vllm.v1.core import single_type_kv_cache_manager

    # 0.29 sizes every KVCacheTensor region from the resolved physical layout,
    # so it must be pinned before get_kv_cache_configs.  Idempotent, so callers
    # that already resolved it (the control plane) are unaffected.
    layout = resolve_diffusion_kv_cache_layout(
        vllm_config,
        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
    )
    single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
    worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
    scheduler_config = generate_scheduler_kv_cache_config(worker_configs)
    # Workers only receive a KVCacheConfig, so the resolved layout travels on it
    # and is adopted into each worker's own CacheConfig before allocation.
    for config in (*worker_configs, scheduler_config):
        config.kv_cache_layout = layout.name
    return worker_configs, scheduler_config


def initialize_diffusion_kv_control_plane(
    executor: DiffusionExecutor,
    od_config: OmniDiffusionConfig,
    *,
    profile_requests: list[OmniDiffusionRequest] | None = None,
    device: torch.device | None = None,
) -> tuple[KVCacheConfig, int, int, VllmConfig] | None:
    """Run the native Worker-spec to Scheduler-config initialization chain."""

    if (
        getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
        is not DiffusionKVCacheMode.PAGED_SCHEDULER
    ):
        return None
    if not profile_requests:
        raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests")
    if device is None:
        device = current_omni_platform.get_torch_device(0)
    vllm_config = create_diffusion_vllm_config(device, od_config)
    worker_specs = executor.get_kv_cache_specs()
    available_memory = executor.determine_available_kv_memory(profile_requests)
    if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus:
        raise ValueError(
            "Diffusion KV initialization rank count mismatch: "
            f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}"
        )
    worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs(
        vllm_config,
        worker_specs,
        available_memory,
    )
    scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
        scheduler_kv_cache_config,
        vllm_config,
    )
    # Native sizing may auto-fit ``max_model_len=-1`` against the profiled
    # memory budget.  Forward that final value so every Worker sizes its
    # physical BlockTables from the same bound as the Scheduler.
    resolved_max_model_len = vllm_config.model_config.max_model_len
    if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0:
        raise ValueError("native Diffusion KV sizing must resolve a positive max_model_len")
    executor.set_kv_cache_configs(worker_configs, resolved_max_model_len)
    return (
        scheduler_kv_cache_config,
        scheduler_block_size,
        hash_block_size,
        vllm_config,
    )