| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import numpy as np |
| 4 | import torch |
| 5 | |
| 6 | from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor |
| 7 | |
| 8 | |
| 9 | class RequestState: |
| 10 | def __init__( |
| 11 | self, |
| 12 | max_num_reqs: int, |
| 13 | max_model_len: int, |
| 14 | max_num_batched_tokens: int, |
| 15 | num_speculative_steps: int, |
| 16 | vocab_size: int, |
| 17 | device: torch.device, |
| 18 | num_prefill_lookahead: int = 1, |
| 19 | ): |
| 20 | self.max_num_reqs = max_num_reqs |
| 21 | self.max_model_len = max_model_len |
| 22 | self.max_num_batched_tokens = max_num_batched_tokens |
| 23 | self.num_speculative_steps = num_speculative_steps |
| 24 | self.vocab_size = vocab_size |
| 25 | self.device = device |
| 26 | |
| 27 | self.req_id_to_index: dict[str, int] = {} |
| 28 | self.index_to_req_id: dict[int, str] = {} |
| 29 | self.free_indices = list(range(max_num_reqs)) |
| 30 | |
| 31 | # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs) |
| 32 | # depending on the configured max_num_reqs and max_model_len. |
| 33 | # To save GPU memory, we use UVA instead of GPU for this tensor. |
| 34 | self.all_token_ids = StagedWriteTensor( |
| 35 | (self.max_num_reqs, self.max_model_len), |
| 36 | dtype=torch.int32, |
| 37 | device=device, |
| 38 | uva_instead_of_gpu=True, |
| 39 | ) |
| 40 | # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len: |
| 41 | # - prompt_len: Number of tokens in the user-provided prompt. |
| 42 | # - prefill_len: Number of tokens passed into the model runner. |
| 43 | # This can include the prompt and additional partial output tokens, |
| 44 | # so prefill_len >= prompt_len. |
| 45 | # Usually, prefill_len equals prompt_len, but in cases such as resumption after |
| 46 | # preemption, prefill_len may be greater. Differentiating between these values |
| 47 | # is crucial, as certain features such as prompt logprobs or frequency penalties |
| 48 | # must treat prompt and output tokens separately. |
| 49 | self.prompt_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) |
| 50 | self.prefill_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) |
| 51 | # total_len = prompt_len + output_len. It grows as the request progresses. |
| 52 | self.total_len = StagedWriteTensor( |
| 53 | self.max_num_reqs, dtype=torch.int32, device=device |
| 54 | ) |
| 55 | |
| 56 | # Number of computed tokens. |
| 57 | self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32) |
| 58 | self.num_computed_tokens = StagedWriteTensor( |
| 59 | self.max_num_reqs, dtype=torch.int32, device=device |
| 60 | ) |
| 61 | # Optimistic CPU mirror of num_computed_tokens (upper bound on GPU value). |
| 62 | self.num_computed_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32) |
| 63 | |
| 64 | # Last sampled tokens. |
| 65 | self.last_sampled_tokens = torch.zeros( |
| 66 | self.max_num_reqs, 1, dtype=torch.int64, device=device |
| 67 | ) |
| 68 | |
| 69 | # Max total seq length (prompt_len + max_tokens). |
| 70 | self.max_seq_len = np.zeros(self.max_num_reqs, dtype=np.int32) |
| 71 | |
| 72 | # Draft tokens. |
| 73 | self.draft_tokens = torch.zeros( |
| 74 | self.max_num_reqs, |
| 75 | self.num_speculative_steps, |
| 76 | dtype=torch.int64, |
| 77 | device=device, |
| 78 | ) |
| 79 | |
| 80 | self.next_prefill_tokens = torch.zeros( |
| 81 | num_prefill_lookahead, self.max_num_reqs, dtype=torch.int32, device=device |
| 82 | ) |
| 83 | |
| 84 | @property |
| 85 | def num_reqs(self) -> int: |
| 86 | return len(self.req_id_to_index) |
| 87 | |
| 88 | def add_request( |
| 89 | self, |
| 90 | req_id: str, |
| 91 | prompt_len: int, |
| 92 | all_token_ids: list[int], |
| 93 | num_computed_tokens: int, |
| 94 | max_tokens: int, |
| 95 | ) -> None: |
| 96 | assert len(self.free_indices) > 0, "No free indices" |
| 97 | req_idx = self.free_indices.pop() |
| 98 | self.req_id_to_index[req_id] = req_idx |
| 99 | self.index_to_req_id[req_idx] = req_id |
| 100 | |
| 101 | self.max_seq_len[req_idx] = prompt_len + max_tokens |
| 102 | self.prompt_len.np[req_idx] = prompt_len |
| 103 | prefill_len = len(all_token_ids) |
| 104 | assert prefill_len >= prompt_len, ( |
| 105 | f"prefill_len {prefill_len} < prompt_len {prompt_len}" |
| 106 | ) |
| 107 | self.prefill_len.np[req_idx] = prefill_len |
| 108 | self.total_len.stage_write_elem(req_idx, prefill_len) |
| 109 | self.all_token_ids.stage_write(req_idx, 0, all_token_ids) |
| 110 | self.num_computed_prefill_tokens[req_idx] = num_computed_tokens |
| 111 | self.num_computed_tokens_np[req_idx] = num_computed_tokens |
| 112 | self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens) |
| 113 | |
| 114 | self.draft_tokens[req_idx].zero_() |
| 115 | |
| 116 | def apply_staged_writes(self) -> None: |
| 117 | self.prompt_len.copy_to_uva() |
| 118 | self.prefill_len.copy_to_uva() |
| 119 | self.total_len.apply_write() |
| 120 | self.all_token_ids.apply_write() |
| 121 | self.num_computed_tokens.apply_write() |
| 122 | |
| 123 | def remove_request(self, req_id: str) -> int | None: |
| 124 | """Return the freed slot index, or None if the request was not found.""" |
| 125 | req_idx = self.req_id_to_index.pop(req_id, None) |
| 126 | if req_idx is None: |
| 127 | return None |
| 128 | self.index_to_req_id.pop(req_idx, None) |
| 129 | self.free_indices.append(req_idx) |
| 130 | return req_idx |