hsliu
51d2570464
- Add comprehensive PRD, architecture design, and test design documents - Implement core modules: OmniLLM, AsyncOmniLLM, stage configurations - Add DiT scheduler and cache manager for diffusion models - Implement CLI integration with --omni flag support - Add API server and plugin system for vLLM integration - Create comprehensive test suite with fixtures - Update dependencies to vLLM 0.10.2 and PyTorch 2.8.0 - Add conda environment setup and package installation - Implement stage-based processing architecture - Add multimodal output processing capabilities This commit establishes the foundation for multi-modality models inference and serving with non-autoregressive structures.
40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
|
|
from vllm_omni.request import OmniRequest
|
|
from typing import List
|
|
from threading import Lock, Condition
|
|
from vllm_omni.config import OmniConfig
|
|
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
from vllm.v1.core.sched import SchedulerInterface
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
|
|
from vllm.v1.structured_output import StructuredOutputManager
|
|
|
|
|
|
class OmniScheduler(SchedulerInterface):
|
|
"""
|
|
OmniScheduler: Scheduler for vLLM-omni multimodal processing.
|
|
|
|
This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive
|
|
processing with additional fields and methods specific to vLLM-omni.
|
|
"""
|
|
def __init__(self,
|
|
omni_config: OmniConfig,
|
|
kv_cache_config: KVCacheConfig,
|
|
structured_output_manager: StructuredOutputManager,
|
|
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
|
|
include_finished_set: bool = False,
|
|
log_stats: bool = False,
|
|
):
|
|
super().__init__(
|
|
vllm_config=omni_config.vllm_config,
|
|
kv_cache_config=kv_cache_config,
|
|
multimodal_registry=mm_registry,
|
|
structured_output_manager=structured_output_manager,
|
|
include_finished_set=include_finished_set,
|
|
log_stats=log_stats,
|
|
)
|
|
self.omni_config = omni_config
|
|
|
|
def schedule(self, requests: List[OmniRequest]) -> List[OmniRequest]:
|
|
# TODO: Implement scheduling logic
|
|
pass |