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.
230 lines
8.6 KiB
Python
230 lines
8.6 KiB
Python
"""
|
|
Diffusion scheduler for DiT (Diffusion Transformer) models in vLLM-omni.
|
|
"""
|
|
|
|
from typing import Dict, List, Optional, Any, TYPE_CHECKING
|
|
from vllm.v1.core.sched.scheduler import SchedulerInterface
|
|
from vllm.config import VllmConfig, KVCacheConfig
|
|
from vllm.v1.core.sched.structured_output_manager import StructuredOutputManager
|
|
from vllm.v1.core.sched.mm_registry import MultiModalRegistry, MULTIMODAL_REGISTRY
|
|
|
|
from ..dit_cache_manager import DiTCacheManager
|
|
from ...config import DiTCacheConfig
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
|
|
|
|
|
class OmniDiffusionScheduler(SchedulerInterface):
|
|
"""Scheduler for DiT models with caching optimization."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
kv_cache_config: KVCacheConfig,
|
|
dit_cache_config: DiTCacheConfig,
|
|
structured_output_manager: StructuredOutputManager,
|
|
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
|
|
include_finished_set: bool = False,
|
|
log_stats: bool = False,
|
|
) -> None:
|
|
super().__init__(
|
|
vllm_config=vllm_config,
|
|
kv_cache_config=kv_cache_config,
|
|
structured_output_manager=structured_output_manager,
|
|
mm_registry=mm_registry,
|
|
include_finished_set=include_finished_set,
|
|
log_stats=log_stats,
|
|
)
|
|
|
|
self.dit_cache_config = dit_cache_config
|
|
self.dit_cache_manager = DiTCacheManager(dit_cache_config)
|
|
|
|
# DiT-specific scheduling state
|
|
self.dit_requests: Dict[str, Dict[str, Any]] = {}
|
|
self.current_step: int = 0
|
|
self.max_steps: int = 50 # Default, will be updated from config
|
|
|
|
def schedule(self) -> "SchedulerOutput":
|
|
"""Schedule DiT requests with caching optimization."""
|
|
# Get pending requests
|
|
pending_requests = self._get_pending_requests()
|
|
|
|
if not pending_requests:
|
|
return self._create_empty_scheduler_output()
|
|
|
|
# Apply DiT-specific scheduling logic
|
|
scheduled_requests = self._schedule_dit_requests(pending_requests)
|
|
|
|
# Create scheduler output
|
|
return self._create_scheduler_output(scheduled_requests)
|
|
|
|
def _get_pending_requests(self) -> List[Dict[str, Any]]:
|
|
"""Get pending DiT requests."""
|
|
# This would integrate with vLLM's request management system
|
|
# For now, we'll return a mock implementation
|
|
return []
|
|
|
|
def _schedule_dit_requests(self, requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""Apply DiT-specific scheduling logic."""
|
|
scheduled_requests = []
|
|
|
|
for request in requests:
|
|
# Check cache for this request
|
|
cached_result = self.dit_cache_manager.get_cache(request['request_id'])
|
|
|
|
if cached_result is not None:
|
|
# Use cached result
|
|
request['cached'] = True
|
|
request['cached_result'] = cached_result
|
|
else:
|
|
# Allocate cache for new request
|
|
cache_tensor = self.dit_cache_manager.allocate_cache(
|
|
request['request_id'],
|
|
request.get('cache_size', 1024)
|
|
)
|
|
request['cache_tensor'] = cache_tensor
|
|
request['cached'] = False
|
|
|
|
# Apply DiT-specific scheduling policies
|
|
request = self._apply_dit_scheduling_policies(request)
|
|
scheduled_requests.append(request)
|
|
|
|
return scheduled_requests
|
|
|
|
def _apply_dit_scheduling_policies(self, request: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Apply DiT-specific scheduling policies."""
|
|
# For now, we'll use a simple FIFO policy
|
|
# In practice, this could include:
|
|
# - Priority-based scheduling
|
|
# - Batch size optimization
|
|
# - Memory-aware scheduling
|
|
# - Quality vs speed trade-offs
|
|
|
|
request['priority'] = request.get('priority', 0)
|
|
request['batch_size'] = request.get('batch_size', 1)
|
|
request['scheduling_policy'] = 'fifo'
|
|
|
|
return request
|
|
|
|
def _create_empty_scheduler_output(self) -> "SchedulerOutput":
|
|
"""Create an empty scheduler output."""
|
|
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
|
|
|
return SchedulerOutput(
|
|
scheduled_seq_groups=[],
|
|
ignored_seq_groups=[],
|
|
preempted_seq_groups=[],
|
|
num_preemption_groups=0,
|
|
num_batched_tokens=0,
|
|
blocks_to_swap_in={},
|
|
blocks_to_swap_out={},
|
|
blocks_to_copy={},
|
|
ignored_seq_groups=[],
|
|
num_ignored_seq_groups=0,
|
|
)
|
|
|
|
def _create_scheduler_output(self, scheduled_requests: List[Dict[str, Any]]) -> "SchedulerOutput":
|
|
"""Create scheduler output from scheduled requests."""
|
|
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
|
|
|
# Convert scheduled requests to vLLM's expected format
|
|
scheduled_seq_groups = []
|
|
for request in scheduled_requests:
|
|
seq_group = self._create_seq_group_from_request(request)
|
|
scheduled_seq_groups.append(seq_group)
|
|
|
|
return SchedulerOutput(
|
|
scheduled_seq_groups=scheduled_seq_groups,
|
|
ignored_seq_groups=[],
|
|
preempted_seq_groups=[],
|
|
num_preemption_groups=0,
|
|
num_batched_tokens=sum(req.get('num_tokens', 0) for req in scheduled_requests),
|
|
blocks_to_swap_in={},
|
|
blocks_to_swap_out={},
|
|
blocks_to_copy={},
|
|
ignored_seq_groups=[],
|
|
num_ignored_seq_groups=0,
|
|
)
|
|
|
|
def _create_seq_group_from_request(self, request: Dict[str, Any]) -> Any:
|
|
"""Create a sequence group from a DiT request."""
|
|
# This would create a proper sequence group
|
|
# For now, we'll return a mock implementation
|
|
from vllm.v1.core.sched.sequence import SequenceGroup
|
|
|
|
# Mock sequence group creation
|
|
# In practice, this would properly create a SequenceGroup
|
|
# with the appropriate metadata for DiT processing
|
|
return None
|
|
|
|
def update_from_output(
|
|
self,
|
|
scheduler_output: "SchedulerOutput",
|
|
model_output: Any,
|
|
) -> List[Any]:
|
|
"""Update scheduler state from model output."""
|
|
# Update cache with new results
|
|
for seq_group in scheduler_output.scheduled_seq_groups:
|
|
if hasattr(seq_group, 'request_id'):
|
|
# Store result in cache
|
|
self.dit_cache_manager.store_cache(
|
|
seq_group.request_id,
|
|
model_output
|
|
)
|
|
|
|
# Update DiT-specific state
|
|
self.current_step += 1
|
|
|
|
# Call parent update method
|
|
return super().update_from_output(scheduler_output, model_output)
|
|
|
|
def add_request(self, request_id: str, **kwargs) -> None:
|
|
"""Add a new DiT request to the scheduler."""
|
|
# Store request metadata
|
|
self.dit_requests[request_id] = {
|
|
'request_id': request_id,
|
|
'added_time': self.current_step,
|
|
'status': 'pending',
|
|
**kwargs
|
|
}
|
|
|
|
# Call parent add_request method
|
|
super().add_request(request_id, **kwargs)
|
|
|
|
def remove_request(self, request_id: str) -> None:
|
|
"""Remove a DiT request from the scheduler."""
|
|
# Clean up request metadata
|
|
if request_id in self.dit_requests:
|
|
del self.dit_requests[request_id]
|
|
|
|
# Release cache for this request
|
|
self.dit_cache_manager.release_cache(request_id)
|
|
|
|
# Call parent remove_request method
|
|
super().remove_request(request_id)
|
|
|
|
def get_dit_request_info(self, request_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Get information about a specific DiT request."""
|
|
return self.dit_requests.get(request_id)
|
|
|
|
def get_cache_statistics(self) -> Dict[str, Any]:
|
|
"""Get cache statistics for monitoring."""
|
|
return self.dit_cache_manager.get_statistics()
|
|
|
|
def clear_expired_cache(self) -> None:
|
|
"""Clear expired cache entries."""
|
|
self.dit_cache_manager.clear_expired_cache()
|
|
|
|
def set_max_steps(self, max_steps: int) -> None:
|
|
"""Set the maximum number of diffusion steps."""
|
|
self.max_steps = max_steps
|
|
|
|
def get_current_step(self) -> int:
|
|
"""Get the current diffusion step."""
|
|
return self.current_step
|
|
|
|
def reset_step_counter(self) -> None:
|
|
"""Reset the diffusion step counter."""
|
|
self.current_step = 0
|