Class CheckpointMemoryPolicy

examples/server/server.py:12857–12960

class CheckpointMemoryPolicy(MemoryPolicy):
    def exact_checkpoint_match(self, tokens: Sequence[int]) -> Tuple[int, int]:
        match_seq_id, match_length = self.scheduler.radix_trie.longest_prefix(
            tokens,
            preferred_sequences=self.scheduler.free_sequences,
            exact_only=True,
        )
        if match_seq_id not in self.scheduler.free_sequences:
            return -1, 0
        return match_seq_id, match_length

    def can_admit(self, request: CompletionRequest) -> bool:
        match_seq_id, match_length = self.exact_checkpoint_match(request.prompt_tokens)
        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)
        prompt_length = len(request.prompt_tokens)
        prompt_kv = request.prompt_plan.eval_token_count
        generation_kv = self.generation_kv_for_request(request, prompt_length)
        if self.try_set_sequence_cache_match(
            request,
            resident_reuse_len=match_length,
            required_sequence_ids=request.internal_completion_count,
            required_attn_kv=self.attention_kv_required(
                prompt_kv,
                reused_kv=0,
                generation_kv=generation_kv,
            ),
            skip_attention_budget_when_unbounded=True,
        ):
            return True
        required_attn_kv = self.attention_kv_required(
            prompt_kv,
            reused_kv=max(0, min(match_length, request.prompt_plan.length)),
            generation_kv=generation_kv,
        )
        if len(self.scheduler.unused_sequences) >= required_sequence_ids and (
            not self.scheduler.model.has_attention_budget
            or self.scheduler.sequence_history.size + required_attn_kv <= self.scheduler.model.n_ctx
        ):
            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 and (
                not self.scheduler.model.has_attention_budget
                or self.scheduler.sequence_history.size + required_attn_kv <= self.scheduler.model.n_ctx
            ):
                request.match_sequence_id, request.match_length = self.exact_checkpoint_match(
                    request.prompt_tokens,
                )
                return True
        return False

    def admit_request(self, request: CompletionRequest) -> None:
        match_seq_id = request.match_sequence_id
        match_length = request.match_length