Determine the number of tokens and KV blocks to reserve for a given request. Given a UID (this UID may not be recognized by the model yet), this will return the number of tokens and blocks to reserve for the request. Arguments: uid (int): The UID of the
(self, uid: int, max_request_tokens: int, max_request_blocks)
| 156 | return logits |
| 157 | |
| 158 | def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[int, torch.Tensor]: |
| 159 | """ |
| 160 | Determine the number of tokens and KV blocks to reserve for a given request. Given a UID |
| 161 | (this UID may not be recognized by the model yet), this will return the number of tokens |
| 162 | and blocks to reserve for the request. |
| 163 | |
| 164 | Arguments: |
| 165 | uid (int): The UID of the sequence (as tracked by the scheduling entity). If |
| 166 | this is a new sequence (with a UID unknown to the inference engine), then |
| 167 | an empty placeholder is created to pass to the occupancy logic. |
| 168 | n_tokens (int): The number of tokens to hypothetically send. |
| 169 | |
| 170 | Returns: |
| 171 | Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks |
| 172 | required to schedule the sequence. |
| 173 | """ |
| 174 | seq_desc = self._state_manager.get_sequence(uid) |
| 175 | if seq_desc is None: |
| 176 | if (self._state_manager.n_tracked_sequences == self._config.state_manager.max_tracked_sequences): |
| 177 | return (0, 0) |
| 178 | seq_desc = PlaceholderSequenceDescriptor() |
| 179 | |
| 180 | req_tokens, req_blocks = self._model.get_kv_requirements(seq_desc, max_request_tokens, max_request_blocks) |
| 181 | |
| 182 | return (req_tokens, req_blocks) |
| 183 | |
| 184 | def can_schedule(self, uids: Iterable[int], lengths: Iterable[int]) -> SchedulingResult: |
| 185 | """ |