Class PartitionedAttentionMemoryPolicy

examples/server/server.py:12813–12854

class PartitionedAttentionMemoryPolicy(AttentionMemoryPolicy):
    def can_admit(self, request: CompletionRequest) -> bool:
        match_seq_id, match_length = self.match_prefix(request.prompt_tokens)
        match_length = request.prompt_plan.clamp_to_reusable_boundary(match_length)
        request.match_sequence_id = match_seq_id
        request.match_length = match_length
        claimable = match_seq_id in self.scheduler.free_sequences
        required_sequence_ids = request.internal_completion_count - int(claimable)
        reuse_len = self.reuse_len_for_request(request, match_length)
        if self.try_set_sequence_cache_match(
            request,
            resident_reuse_len=reuse_len,
            required_sequence_ids=request.internal_completion_count,
        ):
            return True
        if len(self.scheduler.unused_sequences) >= required_sequence_ids:
            return True
        best_free = match_seq_id if claimable else None
        for seq_id in self.reclaim_order(best_free):
            self.scheduler.delete_free_sequence(seq_id)
            if len(self.scheduler.unused_sequences) >= required_sequence_ids:
                request.match_sequence_id, request.match_length = self.match_prefix(
                    request.prompt_tokens,
                )
                request.match_length = request.prompt_plan.clamp_to_reusable_boundary(
                    request.match_length
                )
                return True
        return False

    def copy_prompt_state(
        self,
        source_sequence_id: int,
        dest_sequence_id: int,
        keep_len: int,
    ) -> None:
        self.scheduler.copy_sequence_state(
            source_sequence_id,
            dest_sequence_id,
            keep_len,
            copy_all_state=True,
        )


class CheckpointMemoryPolicy(MemoryPolicy):