8ed5421aaa · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/engine/test_engine_core.py

+10

@@ -102,14 +102,24 @@ def test_engine_core(monkeypatch):

102 engine_core.add_request(req)

103 assert len(engine_core.scheduler.waiting) == 1

104 assert len(engine_core.scheduler.running) == 0

105 + assert engine_core.scheduler.has_unfinished_requests()

106 + assert not engine_core.scheduler.has_finished_requests()

107

108 _ = engine_core.step()

109 assert len(engine_core.scheduler.waiting) == 0

110 assert len(engine_core.scheduler.running) == 1

111 + assert engine_core.scheduler.has_unfinished_requests()

112 + assert not engine_core.scheduler.has_finished_requests()

113

114 engine_core.abort_requests([request_id])

115 assert len(engine_core.scheduler.waiting) == 0

116 assert len(engine_core.scheduler.running) == 0

117 + assert not engine_core.scheduler.has_unfinished_requests()

118 + assert engine_core.scheduler.has_finished_requests()

119 +

120 + _ = engine_core.step()

121 + assert not engine_core.scheduler.has_unfinished_requests()

122 + assert not engine_core.scheduler.has_finished_requests()

123

124 # Add, step, abort 1 of the 3.

125 req0 = make_request()

tests/v1/engine/test_engine_core_client.py

+2 -2

@@ -50,7 +50,7 @@ def loop_until_done(client: EngineCoreClient, outputs: dict):

50 engine_core_outputs = client.get_output().outputs

51

52 if len(engine_core_outputs) == 0:

53 - break

53 + continue

54

55 all_finished = True

56 for out in engine_core_outputs:

@@ -68,7 +68,7 @@ async def loop_until_done_async(client: EngineCoreClient, outputs: dict):

68 engine_core_outputs = (await client.get_output_async()).outputs

69

70 if len(engine_core_outputs) == 0:

71 - break

71 + continue

72

73 all_finished = True

74 for out in engine_core_outputs:

vllm/v1/core/scheduler.py

+10 -1

@@ -682,7 +682,8 @@ class Scheduler:

682 assert RequestStatus.is_finished(finished_status)

683 if isinstance(request_ids, str):

684 request_ids = (request_ids, )

685 - request_ids = set(request_ids)

685 + else:

686 + request_ids = set(request_ids)

687

688 for req_id in request_ids:

689 request = self.requests.get(req_id)

@@ -714,6 +715,14 @@ class Scheduler:

715 def has_unfinished_requests(self) -> bool:

716 return self.get_num_unfinished_requests() > 0

717

718 + def has_finished_requests(self) -> bool:

719 + return len(self.finished_req_ids) > 0

720 +

721 + def has_requests(self):

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

723 + not yet returned in SchedulerOutputs."""

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

725 +

726 def get_num_unscheduled_requests(self) -> int:

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

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

vllm/v1/engine/async_llm.py

+3 -3

@@ -253,13 +253,14 @@ class AsyncLLM(EngineClient):

253 while True:

254 # 1) Pull EngineCoreOutputs from the EngineCore.

255 outputs = await self.engine_core.get_output_async()

256 + num_outputs = len(outputs.outputs)

257

257 - iteration_stats = IterationStats() if self.log_stats else None

258 + iteration_stats = IterationStats() if (

259 + self.log_stats and num_outputs) else None

260

261 # Split outputs into chunks of at most

262 # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the

263 # event loop for too long.

262 - num_outputs = len(outputs.outputs)

264 if num_outputs <= VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:

265 slices = (outputs.outputs, )

266 else:

@@ -313,7 +314,6 @@ class AsyncLLM(EngineClient):

314 return

315

316 assert scheduler_stats is not None

316 - assert iteration_stats is not None

317 for stat_logger in self.stat_loggers:

318 stat_logger.record(scheduler_stats=scheduler_stats,

319 iteration_stats=iteration_stats)

vllm/v1/engine/core.py

+4 -2

@@ -153,7 +153,9 @@ class EngineCore:

153 def step(self) -> EngineCoreOutputs:

154 """Schedule, execute, and make output."""

155

156 - if not self.scheduler.has_unfinished_requests():

156 + # Check for any requests remaining in the scheduler - unfinished,

157 + # or finished and not yet removed from the batch.

158 + if not self.scheduler.has_requests():

159 return EngineCoreOutputs(

160 outputs=[],

161 scheduler_stats=self.scheduler.make_stats(),

@@ -335,7 +337,7 @@ class EngineCoreProc(EngineCore):

337 # Loop until process is sent a SIGINT or SIGTERM

338 while True:

339 # 1) Poll the input queue until there is work to do.

338 - while not self.scheduler.has_unfinished_requests():

340 + while not self.scheduler.has_requests():

341 logger.debug("EngineCore busy loop waiting.")

342 req = self.input_queue.get()

343 self._handle_client_request(*req)

vllm/v1/metrics/loggers.py

+8 -4

@@ -22,7 +22,7 @@ class StatLoggerBase(ABC):

22

23 @abstractmethod

24 def record(self, scheduler_stats: SchedulerStats,

25 - iteration_stats: IterationStats):

25 + iteration_stats: Optional[IterationStats]):

26 ...

27

28 def log(self): # noqa

@@ -56,10 +56,11 @@ class LoggingStatLogger(StatLoggerBase):

56 return float(np.sum(tracked_stats) / (now - self.last_log_time))

57

58 def record(self, scheduler_stats: SchedulerStats,

59 - iteration_stats: IterationStats):

59 + iteration_stats: Optional[IterationStats]):

60 """Log Stats to standard output."""

61

62 - self._track_iteration_stats(iteration_stats)

62 + if iteration_stats:

63 + self._track_iteration_stats(iteration_stats)

64

65 self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats)

66

@@ -319,7 +320,7 @@ class PrometheusStatLogger(StatLoggerBase):

320 info_gauge.set(1)

321

322 def record(self, scheduler_stats: SchedulerStats,

322 - iteration_stats: IterationStats):

323 + iteration_stats: Optional[IterationStats]):

324 """Log to prometheus."""

325 self.gauge_scheduler_running.set(scheduler_stats.num_running_reqs)

326 self.gauge_scheduler_waiting.set(scheduler_stats.num_waiting_reqs)

@@ -331,6 +332,9 @@ class PrometheusStatLogger(StatLoggerBase):

332 self.counter_gpu_prefix_cache_hits.inc(

333 scheduler_stats.prefix_cache_stats.hits)

334

335 + if iteration_stats is None:

336 + return

337 +

338 self.counter_num_preempted_reqs.inc(iteration_stats.num_preempted_reqs)

339 self.counter_prompt_tokens.inc(iteration_stats.num_prompt_tokens)

340 self.counter_generation_tokens.inc(

vllm/v1/outputs.py

+10

@@ -80,3 +80,13 @@ class ModelRunnerOutput:

80 # [prompt_len, num_prompt_logprobs]

81 # [prompt_len]

82 prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]]

83 +

84 +

85 +EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

86 + req_ids=[],

87 + req_id_to_index={},

88 + sampled_token_ids=[],

89 + spec_token_ids=None,

90 + logprobs=None,

91 + prompt_logprobs_dict={},

92 +)

vllm/v1/worker/gpu_model_runner.py

+6 -3

@@ -32,7 +32,8 @@ from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

32 from vllm.v1.engine.mm_input_cache import MMInputCacheClient

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

34 KVCacheSpec)

35 -from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput

35 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,

36 + ModelRunnerOutput)

37 from vllm.v1.sample.metadata import SamplingMetadata

38 from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID, RejectionSampler

39 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

@@ -919,6 +920,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

920 intermediate_tensors: Optional[IntermediateTensors] = None,

921 ) -> Union[ModelRunnerOutput, torch.Tensor]:

922 self._update_states(scheduler_output)

923 + if not scheduler_output.total_num_scheduled_tokens:

924 + # Return empty ModelRunnerOuptut if there's no work to do.

925 + return EMPTY_MODEL_RUNNER_OUTPUT

926

927 if self.is_multimodal_model:

928 # Run the multimodal encoder if any.

@@ -1069,7 +1073,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1073 spec_token_ids = self.generate_draft_token_ids(

1074 valid_sampled_token_ids)

1075

1072 - model_runner_output = ModelRunnerOutput(

1076 + return ModelRunnerOutput(

1077 req_ids=self.input_batch.req_ids,

1078 req_id_to_index=self.input_batch.req_id_to_index,

1079 sampled_token_ids=valid_sampled_token_ids,

@@ -1077,7 +1081,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1081 logprobs=logprobs_lists,

1082 prompt_logprobs_dict=prompt_logprobs_dict,

1083 )

1080 - return model_runner_output

1084

1085 def generate_draft_token_ids(

1086 self,

vllm/v1/worker/tpu_model_runner.py

+5 -1

@@ -29,7 +29,8 @@ from vllm.v1.attention.backends.pallas import (NUM_KV_PAGES_PER_BLOCK,

29 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

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

31 KVCacheSpec)

32 -from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput

32 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,

33 + ModelRunnerOutput)

34 from vllm.v1.utils import bind_kv_cache

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

36

@@ -546,6 +547,9 @@ class TPUModelRunner:

547 ) -> ModelRunnerOutput:

548 # Update cached state

549 self._update_states(scheduler_output)

550 + if not scheduler_output.total_num_scheduled_tokens:

551 + # Return empty ModelRunnerOuptut if there's no work to do.

552 + return EMPTY_MODEL_RUNNER_OUTPUT

553

554 if self.is_multimodal_model:

555 # Run the multimodal encoder if any.

[V1] Eagerly remove finished requests from the batch (#14388)