@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4c21ce9eba

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, 23:30 UTC

[V1] Get input tokens from scheduler (#13339)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Feb 17, 2025 at 11:01 UTC 4c21ce9ebabcdd0892b5b213f04541ad86507a66

4 files changed +139 -139

tests/v1/worker/test_gpu_model_runner.py

@@ -154,6 +154,7 @@ def test_update_states_request_resumed(model_runner):

154 cached_req_data = CachedRequestData(

155 req_id=req_id,

156 resumed_from_preemption=False,

157 + new_token_ids=[],

158 new_block_ids=[],

159 num_computed_tokens=0,

160 )

vllm/v1/core/scheduler.py

+28 -15

@@ -121,6 +121,8 @@ class Scheduler:

121 encoder_budget = self.max_num_encoder_input_tokens

122 # Spec decode-related.

123 scheduled_spec_decode_tokens: Dict[str, List[int]] = {}

124 +

125 + # For logging.

126 scheduled_timestamp = time.monotonic()

127

128 # First, schedule the RUNNING requests.

@@ -187,6 +189,15 @@ class Scheduler:

189 token_budget -= num_new_tokens

190 req_index += 1

191

192 + # Speculative decode related.

193 + if request.spec_token_ids:

194 + num_scheduled_spec_tokens = (num_new_tokens +

195 + request.num_computed_tokens -

196 + request.num_tokens)

197 + if num_scheduled_spec_tokens > 0:

198 + scheduled_spec_decode_tokens[request.request_id] = (

199 + request.spec_token_ids[:num_scheduled_spec_tokens])

200 +

201 # Encoder-related.

202 if encoder_inputs_to_schedule:

203 scheduled_encoder_inputs[request.request_id] = (

@@ -196,11 +207,6 @@ class Scheduler:

207 self.encoder_cache_manager.allocate(request, i)

208 encoder_budget = new_encoder_budget

209

199 - # Speculative decode related.

200 - if request.spec_token_ids:

201 - scheduled_spec_decode_tokens[

202 - request.request_id] = request.spec_token_ids

203 -

210 # Record the LoRAs in scheduled_running_reqs

211 requested_loras: Set[int] = set()

212 if self.lora_config:

@@ -324,23 +330,24 @@ class Scheduler:

330 # Construct the scheduler output.

331 new_reqs_data = [

332 NewRequestData.from_request(req,

327 - req_to_new_block_ids[req.request_id],

328 - req.num_computed_tokens)

333 + req_to_new_block_ids[req.request_id])

334 for req in scheduled_new_reqs

335 ]

336 resumed_reqs_data = [

337 self._make_cached_request_data(

338 req,

339 + num_scheduled_tokens[req.request_id],

340 + len(scheduled_spec_decode_tokens.get(req.request_id, ())),

341 req_to_new_block_ids[req.request_id],

335 - req.num_computed_tokens,

342 resumed_from_preemption=True,

343 ) for req in scheduled_resumed_reqs

344 ]

345 running_reqs_data = [

346 self._make_cached_request_data(

347 req,

348 + num_scheduled_tokens[req.request_id],

349 + len(scheduled_spec_decode_tokens.get(req.request_id, ())),

350 req_to_new_block_ids[req.request_id],

343 - req.num_computed_tokens,

351 resumed_from_preemption=False,

352 ) for req in scheduled_running_reqs

353 ]

@@ -349,8 +356,8 @@ class Scheduler:

356 scheduled_cached_reqs=resumed_reqs_data + running_reqs_data,

357 num_scheduled_tokens=num_scheduled_tokens,

358 total_num_scheduled_tokens=total_num_scheduled_tokens,

352 - scheduled_encoder_inputs=scheduled_encoder_inputs,

359 scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,

360 + scheduled_encoder_inputs=scheduled_encoder_inputs,

361 num_common_prefix_blocks=num_common_prefix_blocks,

362 # finished_req_ids is an existing state in the scheduler,

363 # instead of being newly scheduled in this step.

@@ -366,22 +373,28 @@ class Scheduler:

373 def _make_cached_request_data(

374 self,

375 request: Request,

376 + num_scheduled_tokens: int,

377 + num_scheduled_spec_tokens: int,

378 new_block_ids: List[int],

370 - num_computed_tokens: int,

379 resumed_from_preemption: bool,

380 ) -> "CachedRequestData":

381 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating

382 # them at each scheduling step.

375 - if request.request_id in self._cached_reqs_data:

376 - req_data = self._cached_reqs_data[request.request_id]

383 + num_computed_tokens = request.num_computed_tokens

384 + num_regular_tokens = num_scheduled_tokens - num_scheduled_spec_tokens

385 + new_token_ids = request.all_token_ids[

386 + num_computed_tokens:num_computed_tokens + num_regular_tokens]

387 + req_data = self._cached_reqs_data.get(request.request_id)

388 + if req_data is not None:

389 req_data.resumed_from_preemption = resumed_from_preemption

390 + req_data.new_token_ids = new_token_ids

391 req_data.new_block_ids = new_block_ids

392 req_data.num_computed_tokens = num_computed_tokens

393 else:

394 req_data = CachedRequestData.from_request(request,

395 resumed_from_preemption,

383 - new_block_ids,

384 - num_computed_tokens)

396 + new_token_ids,

397 + new_block_ids)

398 self._cached_reqs_data[request.request_id] = req_data

399 return req_data

400

vllm/v1/core/scheduler_output.py

+8 -7

@@ -30,7 +30,6 @@ class NewRequestData:

30 cls,

31 request: "Request",

32 block_ids: List[int],

33 - num_computed_tokens: int,

33 ) -> "NewRequestData":

34 return cls(

35 req_id=request.request_id,

@@ -41,7 +40,7 @@ class NewRequestData:

40 mm_positions=request.mm_positions,

41 sampling_params=request.sampling_params,

42 block_ids=block_ids,

44 - num_computed_tokens=num_computed_tokens,

43 + num_computed_tokens=request.num_computed_tokens,

44 lora_request=request.lora_request,

45 )

@@ -54,6 +53,7 @@ class CachedRequestData:

53 # the request's block IDs. If True, new_block_ids will be used as the

54 # request's block IDs instead of appending to the existing block IDs.

55 resumed_from_preemption: bool

56 + new_token_ids: List[int]

57 new_block_ids: List[int]

58 num_computed_tokens: int

@@ -62,14 +62,15 @@ class CachedRequestData:

62 cls,

63 request: "Request",

64 resumed_from_preemption: bool,

65 + new_token_ids: List[int],

66 new_block_ids: List[int],

66 - num_computed_tokens: int,

67 ) -> "CachedRequestData":

68 return cls(

69 req_id=request.request_id,

70 resumed_from_preemption=resumed_from_preemption,

71 + new_token_ids=new_token_ids,

72 new_block_ids=new_block_ids,

72 - num_computed_tokens=num_computed_tokens,

73 + num_computed_tokens=request.num_computed_tokens,

74 )

@@ -91,9 +92,9 @@ class SchedulerOutput:

92 # Total number of tokens scheduled for all requests.

93 # Equal to sum(num_scheduled_tokens.values())

94 total_num_scheduled_tokens: int

94 - # req_id -> spec_decode_tokens

95 - # If a request does not have any spec decode tokens, it will

96 - # not be included in the dictionary.

95 + # req_id -> spec_token_ids

96 + # If a request does not have any spec decode tokens, it will not be

97 + # included in the dictionary.

98 scheduled_spec_decode_tokens: Dict[str, List[int]]

99 # req_id -> encoder input indices that need processing.

100 # E.g., if a request has [0, 1], it could mean the vision encoder needs

vllm/v1/worker/gpu_model_runner.py

+102 -117

@@ -2,7 +2,7 @@

3 import gc

4 import time

5 -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, cast

5 +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

7 import numpy as np

8 import torch

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

184 self.max_model_len,

185 self.max_num_tokens),

186 dtype=np.int32)

187 - self.arange_cpu = torch.from_numpy(self.arange_np)

187 # NOTE(woosuk): These tensors are "stateless", i.e., they are literally

188 # a faster version of creating a new tensor every time. Thus, we should

189 # not make any assumptions about the values in these tensors.

@@ -327,7 +326,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):

326 req_state = self.requests[req_id]

327

328 # Update the cached states.

330 - req_state.num_computed_tokens = req_data.num_computed_tokens

329 + num_computed_tokens = req_data.num_computed_tokens

330 + req_state.num_computed_tokens = num_computed_tokens

331 + # Add the sampled token(s) from the previous step (if any).

332 + # This doesn't include "unverified" tokens like spec decode tokens.

333 + num_new_tokens = (num_computed_tokens +

334 + len(req_data.new_token_ids) -

335 + req_state.num_tokens)

336 + new_token_ids = (req_data.new_token_ids[-num_new_tokens:]

337 + if num_new_tokens > 0 else [])

338 + req_state.output_token_ids.extend(new_token_ids)

339 + # Update the block IDs.

340 if not req_data.resumed_from_preemption:

341 # Append the new blocks to the existing block IDs.

342 req_state.block_ids.extend(req_data.new_block_ids)

@@ -346,12 +355,30 @@ class GPUModelRunner(LoRAModelRunnerMixin):

355

356 # Update the persistent batch.

357 self.input_batch.num_computed_tokens_cpu[req_index] = (

349 - req_data.num_computed_tokens)

350 - start_index = len(req_state.block_ids) - len(

351 - req_data.new_block_ids)

358 + num_computed_tokens)

359 + start_index = (len(req_state.block_ids) -

360 + len(req_data.new_block_ids))

361 self.input_batch.block_table.append_row(req_index, start_index,

362 req_data.new_block_ids)

363 + # Add new_token_ids to token_ids_cpu.

364 + start_token_index = num_computed_tokens

365 + end_token_index = num_computed_tokens + len(req_data.new_token_ids)

366 + self.input_batch.token_ids_cpu[

367 + req_index,

368 + start_token_index:end_token_index] = req_data.new_token_ids

369 + # Add spec_token_ids to token_ids_cpu.

370 + spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(

371 + req_id, [])

372 + if spec_token_ids:

373 + start_index = end_token_index

374 + end_token_index += len(spec_token_ids)

375 + self.input_batch.token_ids_cpu[

376 + req_index, start_index:end_token_index] = spec_token_ids

377 + # NOTE(woosuk): `num_tokens` here may include spec decode tokens.

378 + self.input_batch.num_tokens[req_index] = end_token_index

379

380 + # Check if the batch has changed. If not, we can skip copying the

381 + # sampling metadata from CPU to GPU.

382 batch_changed = len(removed_req_indices) > 0 or len(req_ids_to_add) > 0

383

384 # Add the new or resumed requests to the persistent batch.

@@ -374,7 +401,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

401 return batch_changed

402

403 def _prepare_inputs(

377 - self, scheduler_output: "SchedulerOutput"

404 + self,

405 + scheduler_output: "SchedulerOutput",

406 ) -> Tuple[FlashAttentionMetadata, torch.Tensor]:

407 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

408 assert total_num_scheduled_tokens > 0

@@ -387,24 +415,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):

415

416 # Get the number of scheduled tokens for each request.

417 # TODO: The Python loop can be slow. Optimize.

390 - num_scheduled_tokens_list: List[int] = []

418 + num_scheduled_tokens = np.empty(num_reqs, dtype=np.int32)

419 max_num_scheduled_tokens = 0

392 - all_spec_token_ids: List[int] = []

393 - num_spec_tokens_list: List[int] = []

420 for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

421 assert req_id is not None

422 num_tokens = scheduler_output.num_scheduled_tokens[req_id]

397 - num_scheduled_tokens_list.append(num_tokens)

423 + num_scheduled_tokens[i] = num_tokens

424 max_num_scheduled_tokens = max(max_num_scheduled_tokens,

425 num_tokens)

400 - spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(

401 - req_id, [])

402 - all_spec_token_ids.extend(spec_token_ids)

403 - num_spec_tokens_list.append(len(spec_token_ids))

404 -

405 - num_scheduled_tokens: np.ndarray = np.array(num_scheduled_tokens_list,

406 - dtype=np.int32)

407 - assert max_num_scheduled_tokens > 0

426

427 # Get request indices.

428 # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

@@ -441,78 +459,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

459 token_indices = (positions_np +

460 req_indices * self.input_batch.token_ids_cpu.shape[1])

461

444 - use_spec_decode = len(all_spec_token_ids) > 0

445 - if use_spec_decode:

446 -

447 - # 1. Write spec_token_ids to input batch.

448 - # Step 1. Get req indices that perform spec decode and repeat

449 - # the req indices by the number of spec tokens. Note

450 - # for requests that don't perform spec decode, the

451 - # number of spec tokens is 0 and the req index is

452 - # repeated 0 times.

453 - # E.g., num_spec_tokens_list: [3, 0, 2, 0, 1]

454 - # spec_req_indices: [0, 0, 0, 2, 2, 4]

455 - spec_req_indices = np.repeat(self.arange_np[:num_reqs],

456 - num_spec_tokens_list)

457 - # spec_offsets: offsets within each spec token list.

458 - # E.g., [1, 2, 3, 1, 2, 1], TODO: avoid the for loop here

459 - spec_offsets = np.concatenate(

460 - [self.arange_np[1:val + 1] for val in num_spec_tokens_list])

461 - # spec_seq_offsets: offsets within each sequence.

462 - # E.g., num_computed_tokens_cpu: [1, 4, 3, 6, 2]

463 - # after repeating: [1, 1, 1, 3, 3, 2]

464 - # spec_seq_offsets: [1, 1, 1, 3, 3, 2] + [1, 2, 3, 1, 2, 1]

465 - # = [2, 3, 4, 4, 5, 3]

466 - spec_seq_offsets = np.repeat(

467 - self.input_batch.num_computed_tokens_cpu[:num_reqs],

468 - num_spec_tokens_list) + spec_offsets

469 - # cumsums_spec_offsets: [0, 0, 0, 2M, 2M, 4M] + [2, 3, 4, 4, 5, 3]

470 - cumsums_spec_offsets = (

471 - spec_seq_offsets +

472 - spec_req_indices * self.input_batch.token_ids_cpu.shape[1])

473 - cumsums_spec_offsets = torch.from_numpy(cumsums_spec_offsets).to(

474 - torch.int64)

475 - all_spec_token_ids = torch.tensor(all_spec_token_ids,

476 - device="cpu",

477 - dtype=self.input_ids_cpu.dtype)

478 -

479 - # Step 2. Write spec token ids to input_ids_cpu.

480 - self.input_batch.token_ids_cpu_tensor.flatten().scatter_(

481 - 0, cumsums_spec_offsets, all_spec_token_ids)

482 -

483 - # 2. Get spec decode logits indices.

484 - # E.g., num_scheduled_tokens: [4, 100, 3, 100, 2]

485 - # cu_num_tokens: [4, 104, 107, 207, 209]

486 - # num_spec_tokens_list: [3, 0, 2, 0, 1]

487 - # num_sampled_tokens: [4, 1, 3, 1, 2]

488 - # spec_decode_logits_indices:

489 - # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]

490 - num_spec_tokens_np = np.array(num_spec_tokens_list, dtype=np.int32)

491 - num_sampled_tokens = num_spec_tokens_np + 1

492 - # logits_start_loc: [0, 103, 104, 206, 207]

493 - logits_start_loc = cu_num_tokens - num_sampled_tokens

494 - # [0, 103, 104, 206, 207] ->

495 - # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]

496 - logits_start_loc = np.repeat(logits_start_loc, num_sampled_tokens)

497 - # The following three lines:

498 - # [4, 1, 3, 1, 2] -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

499 - # Step 1. [4, 1, 3, 1, 2] -> [4, 5, 8, 9, 11]

500 - cu_num_sampled_tokens = np.cumsum(num_sampled_tokens)

501 - # Step 2. [4, 5, 8, 9, 11] -> [0, 4, 5, 8, 9]

502 - # -> [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]

503 - cumsums_sampled_offsets = np.repeat(

504 - cu_num_sampled_tokens - num_sampled_tokens, num_sampled_tokens)

505 - # Step 3. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

506 - # - [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]

507 - # -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

508 - total_num_sampled_tokens = num_sampled_tokens.sum()

509 - sampled_arange = (self.arange_np[:total_num_sampled_tokens] -

510 - cumsums_sampled_offsets)

511 -

512 - # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] ->

513 - # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]

514 - spec_decode_logits_indices = logits_start_loc + sampled_arange

515 -

462 # NOTE(woosuk): We use torch.index_select instead of np.take here

463 # because torch.index_select is much faster than np.take for large

464 # tensors.

@@ -606,9 +552,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

552 suffix_kv_lens=suffix_kv_lens,

553 )

554

555 + use_spec_decode = len(

556 + scheduler_output.scheduled_spec_decode_tokens) > 0

557 if use_spec_decode:

610 - logits_indices = torch.from_numpy(spec_decode_logits_indices).to(

611 - self.device, non_blocking=True)

558 + logits_indices = self._calc_spec_decode_metadata(

559 + scheduler_output, cu_num_tokens)

560 else:

561 # NOTE(woosuk): Due to chunked prefills, the batch may contain

562 # partial requests. While we should not sample any token

@@ -762,6 +710,53 @@ class GPUModelRunner(LoRAModelRunnerMixin):

710

711 mrope_pos_ptr += completion_part_len

712

713 + def _calc_spec_decode_metadata(

714 + self,

715 + scheduler_output: "SchedulerOutput",

716 + cu_num_tokens: np.ndarray,

717 + ) -> Tuple[torch.Tensor, torch.Tensor]:

718 + # Get the number of spec decode tokens for each request.

719 + num_reqs = self.input_batch.num_reqs

720 + num_spec_decode_tokens = np.empty(num_reqs, dtype=np.int32)

721 + for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

722 + assert req_id is not None

723 + num_spec_decode_tokens[i] = len(

724 + scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))

725 +

726 + # Get spec decode logits indices.

727 + # E.g., num_scheduled_tokens: [4, 100, 3, 100, 2]

728 + # cu_num_tokens: [4, 104, 107, 207, 209]

729 + # num_spec_tokens_list: [3, 0, 2, 0, 1]

730 + # num_sampled_tokens: [4, 1, 3, 1, 2]

731 + # spec_decode_logits_indices:

732 + # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]

733 + num_sampled_tokens = num_spec_decode_tokens + 1

734 + # logits_start_loc: [0, 103, 104, 206, 207]

735 + logits_start_loc = cu_num_tokens - num_sampled_tokens

736 + # [0, 103, 104, 206, 207] ->

737 + # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]

738 + logits_start_loc = np.repeat(logits_start_loc, num_sampled_tokens)

739 + # The following three lines:

740 + # [4, 1, 3, 1, 2] -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

741 + # Step 1. [4, 1, 3, 1, 2] -> [4, 5, 8, 9, 11]

742 + cu_num_sampled_tokens = np.cumsum(num_sampled_tokens)

743 + # Step 2. [4, 5, 8, 9, 11] -> [0, 4, 5, 8, 9]

744 + # -> [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]

745 + cumsums_sampled_offsets = np.repeat(

746 + cu_num_sampled_tokens - num_sampled_tokens, num_sampled_tokens)

747 + # Step 3. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

748 + # - [0, 0, 0, 0, 4, 5, 5, 5, 8, 9, 9]

749 + # -> [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]

750 + total_num_sampled_tokens = num_sampled_tokens.sum()

751 + sampled_arange = (self.arange_np[:total_num_sampled_tokens] -

752 + cumsums_sampled_offsets)

753 +

754 + # [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207] ->

755 + # [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]

756 + spec_decode_logits_indices = logits_start_loc + sampled_arange

757 + return torch.from_numpy(spec_decode_logits_indices).to(

758 + self.device, non_blocking=True)

759 +

760 def _prepare_sampling(

761 self,

762 batch_changed: bool,

@@ -773,7 +768,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

768 for req_id, req in self.requests.items()}

769

770 sampling_metadata = self.input_batch.make_sampling_metadata(

776 - req_id_output_token_ids, req_to_spec_token_ids, not batch_changed)

771 + req_id_output_token_ids,

772 + req_to_spec_token_ids,

773 + skip_copy=not batch_changed)

774 return sampling_metadata

775

776 def _execute_encoder(self, scheduler_output: "SchedulerOutput"):

@@ -960,28 +957,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):

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

958 # the requests one by one. Optimize.

959 num_reqs = self.input_batch.num_reqs

963 - request_seq_lens: List[Tuple[int, CachedRequestState, int]] = []

960 + req_ids: List[str] = []

961 + # Because `input_batch.req_ids` is a list of length `max_num_reqs`,

962 + # we need to stop at `num_reqs`.

963 + # FIXME(woosuk): This is hacky. Refactor.

964 for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

965 assert req_id is not None

966 + req_ids.append(req_id)

967 req_state = self.requests[req_id]

968 seq_len = (req_state.num_computed_tokens +

969 scheduler_output.num_scheduled_tokens[req_id])

969 - if seq_len >= req_state.num_tokens:

970 - request_seq_lens.append((i, req_state, seq_len))

971 - else:

972 - # Ignore the sampled token from the partial request.

970 + if seq_len < req_state.num_tokens:

971 + # Ignore the sampled token.

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

973 generator = self.input_batch.generators.get(i)

974 if generator is not None:

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

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

977

979 - # num_reqs entries should be non-None

980 - assert all(

981 - req_id is not None for req_id in

982 - self.input_batch.req_ids[:num_reqs]), "req_ids contains None"

983 - req_ids = cast(List[str], self.input_batch.req_ids[:num_reqs])

984 -

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

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

980 logprobs_tensors = sampler_output.logprobs_tensors

@@ -994,29 +987,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):

987 scheduler_output,

988 )

989

997 - # Update batch with the valid generated tokens.

990 + # Get the valid generated tokens.

991 sampled_token_ids = sampler_output.sampled_token_ids

992 max_gen_len = sampled_token_ids.shape[-1]

993 if max_gen_len == 1:

994 + # No spec decode tokens.

995 valid_sampled_token_ids = sampled_token_ids.tolist()

1002 - for i, req_state, seq_len in request_seq_lens:

1003 - token_id = valid_sampled_token_ids[i][0]

1004 - self.input_batch.token_ids_cpu[i, seq_len] = token_id

1005 - req_state.output_token_ids.append(token_id)

1006 - self.input_batch.num_tokens[i] += 1

996 else:

997 + # Includes spec decode tokens.

998 valid_mask = sampled_token_ids != INVALID_TOKEN_ID

999 gen_lens = valid_mask.sum(dim=1).tolist()

1000 + # TODO(woosuk): Optimize this.

1001 valid_sampled_token_ids = [

1002 seq.tolist()

1003 for seq in sampled_token_ids[valid_mask].split(gen_lens)

1004 ]

1014 - self.input_batch.num_tokens[:num_reqs] += gen_lens

1015 - for i, req_state, seq_len in request_seq_lens:

1016 - target_slice = slice(seq_len - gen_lens[i] + 1, seq_len + 1)

1017 - self.input_batch.token_ids_cpu[

1018 - i, target_slice] = valid_sampled_token_ids[i]

1019 - req_state.output_token_ids.extend(valid_sampled_token_ids[i])

1005

1006 model_runner_output = ModelRunnerOutput(

1007 req_ids=req_ids,