Signed-off-by: Zijun Huang <zijunhua@usc.edu> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler, SchedulerInterface
|
|
from vllm_omni.diffusion.sched.interface import (
|
|
CachedRequestData,
|
|
DiffusionRequestStatus,
|
|
DiffusionSchedulerOutput,
|
|
KVPrefetchJob,
|
|
NewRequestData,
|
|
SchedulerRequestState,
|
|
StepBatchSamplingParamsKey,
|
|
)
|
|
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
|
|
from vllm_omni.diffusion.sched.sigma_schedule import BASE_SCHEDULE_KEY, DMD2SigmaSchedule
|
|
from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
|
|
|
|
Scheduler = RequestScheduler
|
|
|
|
__all__ = [
|
|
"DiffusionRequestStatus",
|
|
"CachedRequestData",
|
|
"DiffusionSchedulerOutput",
|
|
"KVPrefetchJob",
|
|
"NewRequestData",
|
|
"SchedulerRequestState",
|
|
"BaseScheduler",
|
|
"SchedulerInterface",
|
|
"StepBatchSamplingParamsKey",
|
|
"BASE_SCHEDULE_KEY",
|
|
"DMD2SigmaSchedule",
|
|
"RequestScheduler",
|
|
"StepScheduler",
|
|
"Scheduler",
|
|
]
|