Class AttentionMemoryPolicy

examples/server/server.py:12685–12740

class AttentionMemoryPolicy(MemoryPolicy):
    def match_prefix(self, tokens: Sequence[int]) -> Tuple[int, int]:
        return self.scheduler.radix_trie.longest_prefix(
            tokens,
            preferred_sequences=self.scheduler.free_sequences,
        )

    @staticmethod
    def reuse_len_for_request(request: CompletionRequest, match_length: int) -> int:
        needs_generation = (
            request.payload.max_tokens != 0
            and request.effective_max_len > len(request.prompt_tokens)
        )
        reuse_len = match_length
        if needs_generation and request.prompt_tokens:
            reuse_len = min(reuse_len, len(request.prompt_tokens) - 1)
        return request.prompt_plan.clamp_to_reusable_boundary(reuse_len)

    def admit_request(self, request: CompletionRequest) -> None:
        match_seq_id = request.match_sequence_id
        match_length = request.match_length
        reuse_len = self.reuse_len_for_request(request, match_length)
        claimable = match_seq_id in self.scheduler.free_sequences
        if request.sequence_cache_match is not None:
            base_seq_id = self.scheduler.claim_unused_sequence()
            reuse_len, request.prompt_logits = self.scheduler.hydrate_sequence_cache_match(
                request,
                base_seq_id,
            )
        elif claimable:
            base_seq_id = self.scheduler.claim_free_sequence(match_seq_id)
            if self.scheduler.radix_trie.length(base_seq_id) > reuse_len:
                self.scheduler.truncate_sequence(base_seq_id, reuse_len)
        else:
            base_seq_id = self.scheduler.claim_unused_sequence()
            if reuse_len > 0 and match_seq_id >= 0:
                self.copy_prompt_state(match_seq_id, base_seq_id, reuse_len)
        sibling_count = request.internal_completion_count - 1
        sibling_seq_ids: List[int] = []
        for _ in range(sibling_count):
            seq_id = self.scheduler.claim_unused_sequence()
            sibling_seq_ids.append(seq_id)
        self.scheduler.activate_request(
            request,
            base_seq_id=base_seq_id,
            sibling_seq_ids=sibling_seq_ids,
        )
        request.prompt_cursor = reuse_len
        if request.prompt_cursor == len(request.prompt_tokens):
            request.prompt_done = True
            self.scheduler.maybe_save_sequence_cache(request)
            self.scheduler.start_completions(
                request,
                prompt_output_index=None,
                prompt_logits=request.prompt_logits,
            )