Wang fuyinandHongsheng Liu
72ee535fcd
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
89 lines
3.5 KiB
Python
89 lines
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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.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],
|
|
) -> tuple[list[KVCacheConfig], KVCacheConfig]:
|
|
"""Build rank-local and Scheduler-native configs using vLLM utilities."""
|
|
|
|
from vllm.v1.core import single_type_kv_cache_manager
|
|
|
|
single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
|
|
worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
|
|
return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
|
|
|
|
|
|
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,
|
|
)
|