update config and entrypoints

requirements-dev.txt

@@ -1,29 +0,0 @@
-# Development dependencies
--r requirements.txt
-
-# Testing
-pytest>=7.0.0
-pytest-asyncio>=0.21.0
-pytest-cov>=4.0.0
-pytest-mock>=3.10.0
-
-# Code quality
-black>=23.0.0
-isort>=5.12.0
-flake8>=6.0.0
-mypy>=1.0.0
-pre-commit>=3.0.0
-
-# Documentation
-sphinx>=6.0.0
-sphinx-rtd-theme>=1.2.0
-myst-parser>=1.0.0
-
-# Examples and notebooks
-jupyter>=1.0.0
-matplotlib>=3.5.0
-seaborn>=0.11.0
-
-# Profiling and debugging
-memory-profiler>=0.60.0
-py-spy>=0.3.14

vllm_omni/config/__init__.py

@@ -2,72 +2,21 @@
 Configuration management for vLLM-omni.
 """
 
-import os
-import yaml
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, Any, Optional
+from vllm_omni.dit_cache_interface import DiTCacheConfig
+
+from vllm.config import VllmConfig
 
 
-def load_config(config_path: Optional[str] = None) -> Dict[str, Any]:
+@dataclass
+class OmniConfig:
     """
-    Load configuration from file or use defaults.
-    
-    Args:
-        config_path: Path to configuration file. If None, loads default config.
-        
-    Returns:
-        Configuration dictionary.
+    The configuration for vLLM-omni.
     """
-    if config_path is None:
-        config_path = os.path.join(
-            Path(__file__).parent.parent.parent,
-            "configs",
-            "default.yaml"
-        )
-    
-    config_path = Path(config_path)
-    
-    if not config_path.exists():
-        raise FileNotFoundError(f"Configuration file not found: {config_path}")
-    
-    with open(config_path, 'r', encoding='utf-8') as f:
-        config = yaml.safe_load(f)
-    
-    return config
 
-
-def save_config(config: Dict[str, Any], config_path: str) -> None:
-    """
-    Save configuration to file.
-    
-    Args:
-        config: Configuration dictionary to save.
-        config_path: Path where to save the configuration.
-    """
-    config_path = Path(config_path)
-    config_path.parent.mkdir(parents=True, exist_ok=True)
-    
-    with open(config_path, 'w', encoding='utf-8') as f:
-        yaml.dump(config, f, default_flow_style=False, indent=2)
-
-
-def merge_configs(base_config: Dict[str, Any], override_config: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Merge two configuration dictionaries.
-    
-    Args:
-        base_config: Base configuration.
-        override_config: Configuration to override with.
-        
-    Returns:
-        Merged configuration dictionary.
-    """
-    merged = base_config.copy()
-    
-    for key, value in override_config.items():
-        if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
-            merged[key] = merge_configs(merged[key], value)
-        else:
-            merged[key] = value
-    
-    return merged
+    """vllm config"""
+    vllm_config: VllmConfig = field(default_factory=VllmConfig)
+    """DiT cache config"""
+    dit_cache_config: DiTCacheConfig = field(default_factory=DiTCacheConfig)    # DiT cache config

vllm_omni/dit_cache_interface.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class DiTCacheTensor:
+    """
+    A class for specifying how the workers should initialize the DiT cache.
+    """
+    size: int # the size of the cache tensor in bytes
+
+
+@dataclass
+class DiTCacheConfig:
+    """
+    The DiT cache configuration of a model.
+    """
+    """How should model runner initialize the KV cache tensors for each layer"""
+    dit_cache_tensors: list[DiTCacheTensor]
+    """
+    The DiT cache groups of the model.
+    For models with only one type of DiT, there is only one group that
+    contains all layers.
+    """
+    kv_cache_groups: list[DiTCacheTensor]

vllm_omni/request.py

@@ -1,16 +1,7 @@
 import enum
 import time
 from typing import Optional, Dict, Any
+from vllm.v1.request import Request as vLLMRequest
 
-
-Class Request:
-    # initialize request
-    def __init__(self, request_id: str, request_type: str, request_data: Dict[str, Any]):
-        self.request_id = request_id
-        self.request_type = request_type
-        self.request_data = request_data
-        self.request_time = time.time()
-        self.request_status = RequestStatus.PENDING
-        self.request_result = None
-        self.request_error = None
-        self.request_time = time.time()
+class OmniRequest(vLLMRequest):
+    # pass