@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0c6f5023c3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 18:30 UTC

[V1] Scheduler Refactoring [1/N] - Add Scheduler Interface (#15250)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Cody Yu <hao.yu.cody@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Woosuk Kwon committed Mar 20, 2025 at 17:50 UTC 0c6f5023c390075e842bb7c70bb8f5aa433c584c

17 files changed +182 -45

tests/plugins_tests/test_scheduler_plugins.py

+1 -1

@@ -6,7 +6,7 @@ from vllm.core.scheduler import Scheduler

6 from vllm.engine.arg_utils import EngineArgs

7 from vllm.engine.llm_engine import LLMEngine

8 from vllm.sampling_params import SamplingParams

9 -from vllm.v1.core.scheduler import Scheduler as V1Scheduler

9 +from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler

10 from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine

tests/v1/core/test_scheduler.py

+2 -1

@@ -6,7 +6,8 @@ import pytest

6 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig

7 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

8 from vllm.sampling_params import SamplingParams

9 -from vllm.v1.core.scheduler import Scheduler, SchedulerOutput

9 +from vllm.v1.core.sched.output import SchedulerOutput

10 +from vllm.v1.core.sched.scheduler import Scheduler

11 from vllm.v1.outputs import ModelRunnerOutput

12 from vllm.v1.request import Request, RequestStatus

13 from vllm.v1.structured_output import StructuredOutputManager

tests/v1/worker/test_gpu_model_runner.py

+2 -2

@@ -3,8 +3,8 @@ import pytest

4 from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig

5 from vllm.sampling_params import SamplingParams

6 -from vllm.v1.core.scheduler_output import (CachedRequestData, NewRequestData,

7 - SchedulerOutput)

6 +from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,

7 + SchedulerOutput)

8 from vllm.v1.sample.metadata import SamplingMetadata

9 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

vllm/engine/arg_utils.py

+1 -1

@@ -1695,7 +1695,7 @@ class EngineArgs:

1695 # V1 should use the new scheduler by default.

1696 # Swap it only if this arg is set to the original V0 default

1697 if self.scheduler_cls == EngineArgs.scheduler_cls:

1698 - self.scheduler_cls = "vllm.v1.core.scheduler.Scheduler"

1698 + self.scheduler_cls = "vllm.v1.core.sched.scheduler.Scheduler"

1699

1700 # When no user override, set the default values based on the usage

1701 # context.

vllm/executor/ray_utils.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.utils import get_ip

17 from vllm.worker.worker_base import WorkerWrapperBase

19 if TYPE_CHECKING:

20 - from vllm.v1.core.scheduler import SchedulerOutput

20 + from vllm.v1.core.sched.output import SchedulerOutput

21 from vllm.v1.outputs import ModelRunnerOutput

23 logger = init_logger(__name__)

vllm/v1/attention/backends/flash_attn.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.platforms import current_platform

17 from vllm.utils import cdiv

19 if TYPE_CHECKING:

20 - from vllm.v1.core.scheduler_output import SchedulerOutput

20 + from vllm.v1.core.sched.output import SchedulerOutput

21 from vllm.v1.worker.gpu_input_batch import InputBatch

22 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

vllm/v1/attention/backends/mla/common.py

+1 -1

@@ -212,7 +212,7 @@ except ImportError:

212 from flash_attn import flash_attn_varlen_func

213

214 if TYPE_CHECKING:

215 - from vllm.v1.core.scheduler_output import SchedulerOutput

215 + from vllm.v1.core.sched.output import SchedulerOutput

216 from vllm.v1.worker.gpu_input_batch import InputBatch

217 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

218

vllm/v1/core/sched/__init__.py

vllm/v1/core/sched/interface.py new

+139

@@ -0,0 +1,139 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +from abc import ABC, abstractmethod

3 +from collections.abc import Iterable

4 +from typing import TYPE_CHECKING, Optional, Union

5 +

6 +if TYPE_CHECKING:

7 + from vllm.v1.core.sched.output import SchedulerOutput

8 + from vllm.v1.engine import EngineCoreOutputs

9 + from vllm.v1.metrics.stats import SchedulerStats

10 + from vllm.v1.outputs import ModelRunnerOutput

11 + from vllm.v1.request import Request, RequestStatus

12 +

13 +

14 +class SchedulerInterface(ABC):

15 +

16 + @abstractmethod

17 + def schedule(self) -> "SchedulerOutput":

18 + """Schedule the requests to process in this scheduling step.

19 +

20 + The scheduling decision is made at the iteration level. Each scheduling

21 + step corresponds to a single forward pass of the model. Therefore, this

22 + method is called repeatedly by a busy loop in the engine.

23 +

24 + Essentially, the scheduler produces a dictionary of {req_id: num_tokens}

25 + that specifies how many tokens to process for each request in this

26 + scheduling step. For example, num_tokens can be as large as the number

27 + of prompt tokens for new requests, or it can be 1 for the requests that

28 + are auto-regressively generating new tokens one by one. Otherwise, it

29 + can be somewhere in between in case of chunked prefills, prefix caching,

30 + speculative decoding, etc.

31 +

32 + Additionally, the scheduler also returns useful data about each request

33 + or the batch as a whole. The model runner will use this information in

34 + preparing inputs to the model.

35 +

36 + Returns:

37 + A SchedulerOutput object containing information about the scheduled

38 + requests.

39 + """

40 + raise NotImplementedError

41 +

42 + @abstractmethod

43 + def update_from_output(

44 + self,

45 + scheduler_output: "SchedulerOutput",

46 + model_runner_output: "ModelRunnerOutput",

47 + ) -> "EngineCoreOutputs":

48 + """Update the scheduler state based on the model runner output.

49 +

50 + This method is called after the model runner has processed the scheduled

51 + requests. The model runner output includes generated token ids, draft

52 + token ids for next step, etc. The scheduler uses this information to

53 + update its states, checks the finished requests, and returns the output

54 + for each request.

55 +

56 + Returns:

57 + A EngineCoreOutputs object containing the outputs for each request.

58 + """

59 + raise NotImplementedError

60 +

61 + @abstractmethod

62 + def add_request(self, request: "Request") -> None:

63 + """Add a new request to the scheduler's internal queue.

64 +

65 + Args:

66 + request: The new request being added.

67 + """

68 + raise NotImplementedError

69 +

70 + @abstractmethod

71 + def finish_requests(

72 + self,

73 + request_ids: Union[str, Iterable[str]],

74 + finished_status: "RequestStatus",

75 + ) -> None:

76 + """Finish the requests in the scheduler's internal queue. If the request

77 + is not in the queue, this method will do nothing.

78 +

79 + This method is called in two cases:

80 + 1. When the request is aborted by the client.

81 + 2. When the frontend process detects a stop string of the request after

82 + de-tokenizing its generated tokens.

83 +

84 + Args:

85 + request_ids: A single or a list of request IDs.

86 + finished_status: The finished status of the given requests.

87 + """

88 + raise NotImplementedError

89 +

90 + @abstractmethod

91 + def get_num_unfinished_requests(self) -> int:

92 + """Number of unfinished requests in the scheduler's internal queue."""

93 + raise NotImplementedError

94 +

95 + def has_unfinished_requests(self) -> bool:

96 + """Returns True if there are unfinished requests in the scheduler's

97 + internal queue."""

98 + return self.get_num_unfinished_requests() > 0

99 +

100 + @abstractmethod

101 + def has_finished_requests(self) -> bool:

102 + """Returns True if there are finished requests that need to be cleared.

103 + NOTE: This is different from `not self.has_unfinished_requests()`.

104 +

105 + The scheduler maintains an internal list of the requests finished in the

106 + previous step. This list is returned from the next call to schedule(),

107 + to be sent to the model runner in the next step to clear cached states

108 + for these finished requests.

109 +

110 + This method checks if this internal list of finished requests is

111 + non-empty. This information is useful for DP attention.

112 + """

113 + raise NotImplementedError

114 +

115 + def has_requests(self) -> bool:

116 + """Returns True if there are unfinished requests, or finished requests

117 + not yet returned in SchedulerOutputs."""

118 + return self.has_unfinished_requests() or self.has_finished_requests()

119 +

120 + @abstractmethod

121 + def get_num_unscheduled_requests(self) -> int:

122 + """Number of requests that are not being processed by the executor."""

123 + raise NotImplementedError

124 +

125 + @abstractmethod

126 + def reset_prefix_cache(self) -> bool:

127 + """Reset the prefix cache for KV cache.

128 +

129 + This is particularly required when the model weights are live-updated.

130 + """

131 + raise NotImplementedError

132 +

133 + @abstractmethod

134 + def make_stats(self) -> Optional["SchedulerStats"]:

135 + """Make a SchedulerStats object for logging.

136 +

137 + The SchedulerStats object is created for every scheduling step.

138 + """

139 + raise NotImplementedError

vllm/v1/core/sched/output.py renamed

vllm/v1/core/sched/scheduler.py renamed

+6 -31

@@ -13,8 +13,10 @@ from vllm.logger import init_logger

13 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,

14 compute_encoder_budget)

15 from vllm.v1.core.kv_cache_manager import KVCacheManager

16 -from vllm.v1.core.scheduler_output import (CachedRequestData, NewRequestData,

17 - SchedulerOutput)

16 +from vllm.v1.core.sched.interface import SchedulerInterface

17 +from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,

18 + SchedulerOutput)

19 +from vllm.v1.core.sched.utils import check_stop

20 from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,

21 EngineCoreOutputs)

22 from vllm.v1.metrics.stats import SchedulerStats

@@ -25,7 +27,7 @@ from vllm.v1.structured_output import StructuredOutputManager

27 logger = init_logger(__name__)

28 -class Scheduler:

30 +class Scheduler(SchedulerInterface):

32 def __init__(

33 self,

@@ -602,7 +604,7 @@ class Scheduler:

604

605 # Check for stop and update request state.

606 # This must be called before we make the EngineCoreOutput.

605 - stopped = self._check_stop(request)

607 + stopped = check_stop(request, self.max_model_len)

608 if stopped:

609 self._free_request(request)

610 break

@@ -648,25 +650,6 @@ class Scheduler:

650 scheduler_stats=self.make_stats(),

651 )

652

651 - def _check_stop(self, request: Request) -> bool:

652 - if (request.num_tokens >= self.max_model_len

653 - or request.num_output_tokens >= request.max_tokens):

654 - request.status = RequestStatus.FINISHED_LENGTH_CAPPED

655 - return True

656 -

657 - sampling_params = request.sampling_params

658 - last_token_id = request.output_token_ids[-1]

659 - if (not sampling_params.ignore_eos

660 - and last_token_id == request.eos_token_id):

661 - request.status = RequestStatus.FINISHED_STOPPED

662 - return True

663 -

664 - if last_token_id in (sampling_params.stop_token_ids or ()):

665 - request.status = RequestStatus.FINISHED_STOPPED

666 - request.stop_reason = last_token_id

667 - return True

668 - return False

669 -

653 def add_request(self, request: Request) -> None:

654 self.waiting.append(request)

655 self.requests[request.request_id] = request

@@ -715,17 +698,9 @@ class Scheduler:

698 def get_num_unfinished_requests(self) -> int:

699 return len(self.waiting) + len(self.running)

700

718 - def has_unfinished_requests(self) -> bool:

719 - return self.get_num_unfinished_requests() > 0

720 -

701 def has_finished_requests(self) -> bool:

702 return len(self.finished_req_ids) > 0

703

724 - def has_requests(self):

725 - """Returns True if there are unfinished requests, or finished requests

726 - not yet returned in SchedulerOutputs."""

727 - return self.has_unfinished_requests() or self.has_finished_requests()

728 -

704 def get_num_unscheduled_requests(self) -> int:

705 """Number of requests that are not being processed by the executor."""

706 return self.get_num_unfinished_requests() - len(self.scheduled_req_ids)

vllm/v1/core/sched/utils.py new

+22

@@ -0,0 +1,22 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +from vllm.v1.request import Request, RequestStatus

3 +

4 +

5 +def check_stop(request: Request, max_model_len: int) -> bool:

6 + if (request.num_tokens >= max_model_len

7 + or request.num_output_tokens >= request.max_tokens):

8 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED

9 + return True

10 +

11 + sampling_params = request.sampling_params

12 + last_token_id = request.output_token_ids[-1]

13 + if (not sampling_params.ignore_eos

14 + and last_token_id == request.eos_token_id):

15 + request.status = RequestStatus.FINISHED_STOPPED

16 + return True

17 +

18 + if last_token_id in (sampling_params.stop_token_ids or ()):

19 + request.status = RequestStatus.FINISHED_STOPPED

20 + request.stop_reason = last_token_id

21 + return True

22 + return False

vllm/v1/engine/core.py

+2 -2

@@ -22,8 +22,8 @@ from vllm.transformers_utils.config import (

22 from vllm.utils import (get_exception_traceback, resolve_obj_by_qualname,

23 zmq_socket_ctx)

24 from vllm.v1.core.kv_cache_utils import get_kv_cache_configs

25 -from vllm.v1.core.scheduler import Scheduler as V1Scheduler

26 -from vllm.v1.core.scheduler import SchedulerOutput

25 +from vllm.v1.core.sched.output import SchedulerOutput

26 +from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler

27 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,

28 EngineCoreRequestType, UtilityOutput)

29 from vllm.v1.engine.mm_input_cache import MMInputCacheServer

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -45,7 +45,7 @@ from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

45 if TYPE_CHECKING:

46 import xgrammar as xgr

48 - from vllm.v1.core.scheduler_output import SchedulerOutput

48 + from vllm.v1.core.sched.output import SchedulerOutput

49 else:

50 xgr = LazyLoader("xgr", globals(), "xgrammar")

vllm/v1/worker/gpu_worker.py

+1 -1

@@ -28,7 +28,7 @@ from vllm.v1.worker.worker_base import WorkerBase

28 logger = init_logger(__name__)

30 if TYPE_CHECKING:

31 - from vllm.v1.core.scheduler_output import SchedulerOutput

31 + from vllm.v1.core.sched.output import SchedulerOutput

34 class Worker(WorkerBase):

vllm/v1/worker/tpu_model_runner.py

+1 -1

@@ -37,7 +37,7 @@ from vllm.v1.utils import bind_kv_cache

37 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

39 if TYPE_CHECKING:

40 - from vllm.v1.core.scheduler import SchedulerOutput

40 + from vllm.v1.core.sched.output import SchedulerOutput

42 logger = init_logger(__name__)

vllm/v1/worker/tpu_worker.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.distributed import (ensure_model_parallel_initialized,

17 from vllm.logger import init_logger

18 from vllm.model_executor import set_random_seed

19 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE

20 -from vllm.v1.core.scheduler import SchedulerOutput

20 +from vllm.v1.core.sched.output import SchedulerOutput

21 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

22 KVCacheSpec)

23 from vllm.v1.outputs import ModelRunnerOutput