(self, items: List[CompletionScheduler.BatchItem])
| 14003 | ) |
| 14004 | |
| 14005 | def process_batch(self, items: List[CompletionScheduler.BatchItem]) -> None: |
| 14006 | output_count = sum( |
| 14007 | output_index is not None |
| 14008 | for item in items |
| 14009 | for output_index in item.output_indices |
| 14010 | ) |
| 14011 | for item in items: |
| 14012 | request = self.requests[item.request_id] |
| 14013 | if request.cancelled: |
| 14014 | continue |
| 14015 | self.radix_trie.extend(item.seq_id, item.identity_tokens) |
| 14016 | self.sequence_history.extend( |
| 14017 | item.seq_id, |
| 14018 | item.identity_tokens, |
| 14019 | item.position_increments, |
| 14020 | ) |
| 14021 | if item.kind == "prefill": |
| 14022 | request.capture_prompt_logprobs( |
| 14023 | model=self.model, |
| 14024 | formatter=self.formatter, |
| 14025 | output_indices=item.output_indices, |
| 14026 | output_positions=item.output_positions, |
| 14027 | output_count=output_count, |
| 14028 | output_index_to_logits_index=self.output_index_to_logits_index, |
| 14029 | ) |
| 14030 | prompt_output_index = self.output_index_to_logits_index( |
| 14031 | self.last_output_index(item.output_indices), |
| 14032 | output_count, |
| 14033 | ) |
| 14034 | prompt_logits = ( |
| 14035 | self.model.logits(prompt_output_index) |
| 14036 | if prompt_output_index is not None |
| 14037 | else None |
| 14038 | ) |
| 14039 | prefill = item.require_prefill() |
| 14040 | if prefill.prompt_advance_to is not None: |
| 14041 | request.prompt_cursor = prefill.prompt_advance_to |
| 14042 | else: |
| 14043 | request.prompt_cursor += len(item.identity_tokens) |
| 14044 | if request.prompt_cursor == len(request.prompt_tokens): |
| 14045 | request.prompt_done = True |
| 14046 | request.prompt_logits = prompt_logits |
| 14047 | self.maybe_save_prompt_checkpoint(request) |
| 14048 | self.maybe_save_sequence_cache(request) |
| 14049 | self.start_completions( |
| 14050 | request, |
| 14051 | prompt_output_index=prompt_output_index, |
| 14052 | prompt_logits=request.prompt_logits, |
| 14053 | ) |
| 14054 | else: |
| 14055 | decode = item.require_decode() |
| 14056 | completion = request.completions[decode.completion_index] |
| 14057 | self.process_generation_item( |
| 14058 | completion, |
| 14059 | item, |
| 14060 | output_count, |
| 14061 | ) |
| 14062 | self.finalize_request_if_ready(request) |
no test coverage detected