(self, request: CompletionRequest)
| 12742 | |
| 12743 | class UnifiedAttentionMemoryPolicy(AttentionMemoryPolicy): |
| 12744 | def can_admit(self, request: CompletionRequest) -> bool: |
| 12745 | match_seq_id, match_length = self.match_prefix(request.prompt_tokens) |
| 12746 | match_length = request.prompt_plan.clamp_to_reusable_boundary(match_length) |
| 12747 | request.match_sequence_id = match_seq_id |
| 12748 | request.match_length = match_length |
| 12749 | claimable = match_seq_id in self.scheduler.free_sequences |
| 12750 | required_sequence_ids = request.internal_completion_count - int(claimable) |
| 12751 | prompt_length = len(request.prompt_tokens) |
| 12752 | prompt_kv = request.prompt_plan.eval_token_count |
| 12753 | reuse_len = self.reuse_len_for_request(request, match_length) |
| 12754 | prefix_credit = match_length if claimable else reuse_len |
| 12755 | prefix_credit_kv = max(0, min(prefix_credit, request.prompt_plan.length)) |
| 12756 | generation_kv = self.generation_kv_for_request(request, prompt_length) |
| 12757 | if self.try_set_sequence_cache_match( |
| 12758 | request, |
| 12759 | resident_reuse_len=reuse_len, |
| 12760 | required_sequence_ids=request.internal_completion_count, |
| 12761 | required_attn_kv=self.attention_kv_required( |
| 12762 | prompt_kv, |
| 12763 | reused_kv=0, |
| 12764 | generation_kv=generation_kv, |
| 12765 | ), |
| 12766 | ): |
| 12767 | return True |
| 12768 | required_kv = self.attention_kv_required( |
| 12769 | prompt_kv, |
| 12770 | reused_kv=prefix_credit_kv, |
| 12771 | generation_kv=generation_kv, |
| 12772 | ) |
| 12773 | if ( |
| 12774 | len(self.scheduler.unused_sequences) >= required_sequence_ids |
| 12775 | and self.scheduler.sequence_history.size + required_kv <= self.scheduler.model.n_ctx |
| 12776 | ): |
| 12777 | return True |
| 12778 | best_free = match_seq_id if claimable else None |
| 12779 | for seq_id in self.reclaim_order(best_free): |
| 12780 | if len(self.scheduler.unused_sequences) < required_sequence_ids: |
| 12781 | self.scheduler.delete_free_sequence(seq_id) |
| 12782 | elif self.scheduler.sequence_history.size + required_kv > self.scheduler.model.n_ctx: |
| 12783 | if seq_id == best_free and request.match_length > 0: |
| 12784 | self.scheduler.truncate_free_sequence(seq_id, request.match_length) |
| 12785 | else: |
| 12786 | self.scheduler.delete_free_sequence(seq_id) |
| 12787 | if ( |
| 12788 | len(self.scheduler.unused_sequences) >= required_sequence_ids |
| 12789 | and self.scheduler.sequence_history.size + required_kv <= self.scheduler.model.n_ctx |
| 12790 | ): |
| 12791 | request.match_sequence_id, request.match_length = self.match_prefix( |
| 12792 | request.prompt_tokens, |
| 12793 | ) |
| 12794 | request.match_length = request.prompt_plan.clamp_to_reusable_boundary( |
| 12795 | request.match_length |
| 12796 | ) |
| 12797 | return True |
| 12798 | return False |
| 12799 | |
| 12800 | def copy_prompt_state( |
| 12801 | self, |
nothing calls this directly
no test coverage detected