""" 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