@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 54aa619459

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, 17:15 UTC

[V1] Refactor num_computed_tokens logic (#15307)

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

Cody Yu committed Mar 26, 2025 at 21:54 UTC 54aa619459563c4714f2ac001881dd1b5e3e1d4b

5 files changed +106 -57

tests/v1/core/test_scheduler.py

+13 -3

@@ -244,7 +244,9 @@ def test_schedule_partial_requests():

244 model_runner_output = ModelRunnerOutput(

245 req_ids=[request.request_id for request in requests],

246 req_id_to_index=req_to_index,

247 - sampled_token_ids=[[0] for _ in range(len(requests))],

247 + # Only the first request has a sampled token id because

248 + # the rest requests are still being prefilled.

249 + sampled_token_ids=[[0], [], []],

250 spec_token_ids=None,

251 logprobs=None,

252 prompt_logprobs_dict={},

@@ -266,7 +268,7 @@ def test_schedule_partial_requests():

268

269

270 @pytest.mark.parametrize("enable_prefix_caching", [True, False])

269 -def test_schedule_concurrent_partial_requestse(enable_prefix_caching: bool):

271 +def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):

272 """Test scheduling behavior with concurrent partial requests.

273

274 This test verifies that: there are multiple long prefill requests in the

@@ -304,7 +306,7 @@ def test_schedule_concurrent_partial_requestse(enable_prefix_caching: bool):

306 model_runner_output = ModelRunnerOutput(

307 req_ids=[request.request_id for request in requests],

308 req_id_to_index=req_to_index,

307 - sampled_token_ids=[[0] for _ in range(len(requests))],

309 + sampled_token_ids=[[] for _ in range(len(requests))],

310 spec_token_ids=None,

311 logprobs=None,

312 prompt_logprobs_dict={},

@@ -325,6 +327,14 @@ def test_schedule_concurrent_partial_requestse(enable_prefix_caching: bool):

327 # Schedule the third step. All three requests are running.

328 # First and second requests are in the decode stage.

329 # All the remaining tokens in the third request are processed.

330 + model_runner_output = ModelRunnerOutput(

331 + req_ids=[request.request_id for request in requests],

332 + req_id_to_index=req_to_index,

333 + sampled_token_ids=[[0], [0]] + [[] for _ in range(len(requests) - 2)],

334 + spec_token_ids=None,

335 + logprobs=None,

336 + prompt_logprobs_dict={},

337 + )

338 scheduler.update_from_output(output1, model_runner_output)

339 output2 = scheduler.schedule()

340 assert len(scheduler.running) == 3

tests/v1/engine/test_engine_core.py

+11 -7

@@ -231,8 +231,10 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

231 Test that the engine can handle multiple concurrent batches.

232 """

233

234 - def make_request_with_max_tokens(max_tokens: int) -> EngineCoreRequest:

234 + def make_request_with_max_tokens(req_id: int,

235 + max_tokens: int) -> EngineCoreRequest:

236 request = make_request()

237 + request.request_id = req_id

238 request.sampling_params.max_tokens = max_tokens

239 return request

240

@@ -279,6 +281,8 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

281 # Avoid all requests being scheduled once.

282 enable_prefix_caching=False,

283 max_num_batched_tokens=10,

284 + # Reduce startup time.

285 + enforce_eager=True,

286 )

287 vllm_config = engine_args.create_engine_config()

288 engine_core = EngineCore(vllm_config=vllm_config,

@@ -286,13 +290,13 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

290 executor_class=DummyExecutor)

291 assert engine_core.batch_queue is not None

292

289 - # Add two requests in a row.

290 - req = make_request_with_max_tokens(5)

291 - engine_core.add_request(req)

292 - req = make_request_with_max_tokens(5)

293 - engine_core.add_request(req)

293 + # Add two requests in a row. Each request have 12 prompt tokens.

294 + req0 = make_request_with_max_tokens(0, 5)

295 + engine_core.add_request(req0)

296 + req1 = make_request_with_max_tokens(1, 5)

297 + engine_core.add_request(req1)

298

295 - # First saturate the batch queue.

299 + # Schedule Batch 1: (10, req0)

300 assert engine_core.step_with_batch_queue() is None

301 assert engine_core.batch_queue.qsize() == 1

302 assert engine_core.step_with_batch_queue() is None

vllm/v1/core/sched/scheduler.py

+48 -43

@@ -153,9 +153,9 @@ class Scheduler(SchedulerInterface):

153

154 num_new_tokens = (request.num_tokens_with_spec -

155 request.num_computed_tokens)

156 - if self.scheduler_config.long_prefill_token_threshold > 0:

157 - num_new_tokens = min(

158 - num_new_tokens,

156 + if (0 < self.scheduler_config.long_prefill_token_threshold <

157 + num_new_tokens):

158 + num_new_tokens = (

159 self.scheduler_config.long_prefill_token_threshold)

160 num_new_tokens = min(num_new_tokens, token_budget)

161 assert num_new_tokens > 0

@@ -303,9 +303,9 @@ class Scheduler(SchedulerInterface):

303 num_computed_tokens -= self.block_size

304 num_new_tokens = self.block_size

305 computed_blocks.pop()

306 - if self.scheduler_config.long_prefill_token_threshold > 0:

307 - num_new_tokens = min(

308 - num_new_tokens,

306 + if (0 < self.scheduler_config.long_prefill_token_threshold <

307 + num_new_tokens):

308 + num_new_tokens = (

309 self.scheduler_config.long_prefill_token_threshold)

310 num_new_tokens = min(num_new_tokens, token_budget)

311 assert num_new_tokens > 0

@@ -433,6 +433,18 @@ class Scheduler(SchedulerInterface):

433 grammar_bitmask=grammar_bitmask,

434 )

435

436 + # Advance the number of computed tokens for the request AFTER

437 + # the request is scheduled.

438 + # 1. The scheduler_output of the current step has to include the

439 + # original number of scheduled tokens to determine input IDs.

440 + # 2. Advance the number of computed tokens here allowing us to

441 + # schedule the prefill request again immediately in the next

442 + # scheduling step.

443 + # 3. If some tokens (e.g. spec tokens) are rejected later, the number of

444 + # computed tokens will be adjusted in update_from_output.

445 + for req_id, num_scheduled_token in num_scheduled_tokens.items():

446 + self.requests[req_id].num_computed_tokens += num_scheduled_token

447 +

448 self.finished_req_ids = set()

449 return scheduler_output

450

@@ -561,28 +573,19 @@ class Scheduler(SchedulerInterface):

573

574 req_index = model_runner_output.req_id_to_index[req_id]

575 generated_token_ids = sampled_token_ids[req_index]

564 - if req_id not in scheduler_output.scheduled_spec_decode_tokens:

565 - # When the request's num_computed_tokens catches up

566 - # its num_tokens, the request generates output tokens.

567 - # Otherwise, we ignore the sampler output for the request.

568 - request.num_computed_tokens += num_tokens_scheduled

569 - assert request.num_computed_tokens <= request.num_tokens

570 - else:

571 - # num_computed_tokens_step represents the number of tokens

576 +

577 + scheduled_spec_token_ids = (

578 + scheduler_output.scheduled_spec_decode_tokens.get(req_id))

579 + if scheduled_spec_token_ids:

580 + # num_computed_tokens represents the number of tokens

581 # processed in the current step, considering scheduled

573 - # tokens and rejections.

574 - # It is calculated as:

575 - # num_computed_tokens_step = num_scheduled_tokens -

576 - # num_tokens_rejected,

577 - # where num_tokens_rejected is given by:

582 + # tokens and rejections. If some tokens are rejected,

583 + # num_computed_tokens is decreased by the number of rejected

584 + # tokens, where is given by:

585 # len(scheduled_spec_token_ids) + 1 - len(generated_token_ids).

579 - scheduled_spec_token_ids = (

580 - scheduler_output.scheduled_spec_decode_tokens[req_id])

581 -

582 - num_computed_tokens_step = num_scheduled_tokens[req_id] - (

583 - len(scheduled_spec_token_ids) + 1 -

584 - len(generated_token_ids))

585 - request.num_computed_tokens += num_computed_tokens_step

586 + num_tokens_rejected = (len(scheduled_spec_token_ids) + 1 -

587 + len(generated_token_ids))

588 + request.num_computed_tokens -= num_tokens_rejected

589

590 cached_encoder_input_ids = (

591 self.encoder_cache_manager.get_cached_input_ids(request))

@@ -605,24 +608,26 @@ class Scheduler(SchedulerInterface):

608 new_logprobs = None

609 new_token_ids: list[int] = []

610

608 - if request.num_computed_tokens >= request.num_tokens:

609 - for output_token_id in generated_token_ids:

610 - request.append_output_token_ids(output_token_id)

611 - new_token_ids.append(output_token_id)

612 -

613 - # Check for stop and update request state.

614 - # This must be called before we make the EngineCoreOutput.

615 - stopped = check_stop(request, self.max_model_len)

616 - if stopped:

617 - self._free_request(request)

618 - break

611 + # Append generated tokens and check for stop. Note that if

612 + # a request is still being prefilled, we expect the model runner

613 + # to return empty token ids for the request.

614 + for output_token_id in generated_token_ids:

615 + request.append_output_token_ids(output_token_id)

616 + new_token_ids.append(output_token_id)

617 +

618 + # Check for stop and update request state.

619 + # This must be called before we make the EngineCoreOutput.

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

621 + if stopped:

622 + self._free_request(request)

623 + break

624

620 - # Extract sample logprobs if needed.

621 - if request.sampling_params.logprobs is not None:

622 - assert logprobs is not None

623 - # NOTE: once we support N tokens per step (spec decode),

624 - # the outer lists can be of length > 1.

625 - new_logprobs = logprobs.slice(req_index, req_index + 1)

625 + # Extract sample logprobs if needed.

626 + if (request.sampling_params.logprobs is not None

627 + and logprobs is not None):

628 + # NOTE: once we support N tokens per step (spec decode),

629 + # the outer lists can be of length > 1.

630 + new_logprobs = logprobs.slice(req_index, req_index + 1)

631

632 if new_token_ids and request.use_structured_output:

633 # NOTE: structured_output_request

vllm/v1/sample/rejection_sampler.py

+19

@@ -107,14 +107,33 @@ class RejectionSampler(nn.Module):

107 @staticmethod

108 def parse_output(

109 output_token_ids: torch.Tensor,

110 + ignored_req_idxs: list[int],

111 vocab_size: int,

112 ) -> list[list[int]]:

113 + """Parse the output of the rejection sampler.

114 +

115 + Args:

116 + output_token_ids: The sampled token IDs in shape

117 + [batch_size, max_spec_len + 1]. The rejected tokens are

118 + replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler

119 + and will be filtered out in this function.

120 + ignored_req_idxs: The indices of the requests that should not be

121 + sampled. This is usually because the request is still in the

122 + prefill phase.

123 + vocab_size: The size of the vocabulary.

124 +

125 + Returns:

126 + A list of lists of token IDs.

127 + """

128 output_token_ids_np = output_token_ids.cpu().numpy()

129 # Create mask for valid tokens.

130 valid_mask = ((output_token_ids_np != PLACEHOLDER_TOKEN_ID) &

131 (output_token_ids_np < vocab_size))

132 +

133 + ignored_req_idx_set = set(ignored_req_idxs)

134 outputs = [

135 row[valid_mask[i]].tolist()

136 + if i not in ignored_req_idx_set else []

137 for i, row in enumerate(output_token_ids_np)

138 ]

139 return outputs

vllm/v1/worker/gpu_model_runner.py

+15 -4

@@ -1085,8 +1085,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1085

1086 # TODO(woosuk): The following loop can be slow since it iterates over

1087 # the requests one by one. Optimize.

1088 - for i, generator in self.input_batch.generators.items():

1089 - req_id = self.input_batch.req_ids[i]

1088 + discard_sampled_tokens_req_indices = []

1089 + for i, req_id in enumerate(self.input_batch.req_ids):

1090 req_state = self.requests[req_id]

1091 seq_len = (req_state.num_computed_tokens +

1092 scheduler_output.num_scheduled_tokens[req_id])

@@ -1094,7 +1094,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1094 # Ignore the sampled token for partial prefills.

1095 # Rewind the generator state as if the token was not sampled.

1096 # This relies on cuda-specific torch-internal impl details

1097 - generator.set_offset(generator.get_offset() - 4)

1097 + generator = self.input_batch.generators.get(i)

1098 + if generator is not None:

1099 + generator.set_offset(generator.get_offset() - 4)

1100 + # Record the index of the request that should not be sampled,

1101 + # so that we could clear the sampled tokens before returning.

1102 + discard_sampled_tokens_req_indices.append(i)

1103

1104 # NOTE: GPU -> CPU Sync happens here.

1105 # Move as many CPU operations as possible before this sync point.

@@ -1114,10 +1119,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1119 if max_gen_len == 1:

1120 # No spec decode tokens.

1121 valid_sampled_token_ids = sampled_token_ids.tolist()

1122 + # Mask out the sampled tokens that should not be sampled.

1123 + for i in discard_sampled_tokens_req_indices:

1124 + valid_sampled_token_ids[i].clear()

1125 else:

1126 # Includes spec decode tokens.

1127 valid_sampled_token_ids = self.rejection_sampler.parse_output(

1120 - sampled_token_ids, self.input_batch.vocab_size)

1128 + sampled_token_ids,

1129 + discard_sampled_tokens_req_indices,

1130 + self.input_batch.vocab_size,

1131 + )

1132

1133 if not self.use_spec_decode:

1134 spec_token_ids = None