Extend an LLM Inference Server with a Prefix KV Cache

Context

Candidates for the Anthropic Fellows program report a first online assessment (CodeSignal, about 90 minutes, no AI tools) in which you extend an existing LLM inference server codebase over several parts. Each part unlocks the next, and hidden tests check both outputs and how much prefill work you saved. The version below is a self-contained reconstruction for practice; the real codebase is larger, but the skills are the same: read unfamiliar code fast, use the internal APIs you are given, and keep outputs identical while doing less work.

Starter Code

class KVState:
    """Per-request attention cache. Real code holds K/V tensors per layer; a token list stands in."""
    def __init__(self, tokens=()):
        self.tokens = list(tokens)

    @property
    def length(self): return len(self.tokens)

    def truncate(self, n):
        """New KVState with only the first n positions. Does not modify self."""
        return KVState(self.tokens[:n])

    def clone(self): return KVState(self.tokens)


class Model:
    def __init__(self):
        self.prefill_tokens = 0          # cost counter the tests read

    def prefill(self, tokens, past=None):
        """Run tokens through the model, continuing from `past` if given.
        Returns (logits_of_last_position, new_kv). Needs len(tokens) >= 1. Does not modify `past`."""

    def decode(self, token, kv):
        """One decode step. APPENDS token to kv IN PLACE and returns (logits, kv)."""


def argmax(xs): return max(range(len(xs)), key=xs.__getitem__)


class Server:
    def __init__(self, model):
        self.model = model

    def generate(self, prompt, max_new_tokens):
        logits, kv = self.model.prefill(prompt)
        out = []
        for _ in range(max_new_tokens):
            tok = argmax(logits)
            out.append(tok)
            logits, kv = self.model.decode(tok, kv)
        return out

Server.generate always prefills the whole prompt, even when the same system prompt arrives thousands of times. Make it cheaper without changing any output token.

Parts

1. Exact-match cache. If a prompt was served before, do not prefill it again.
2. Prefix reuse. If a new prompt starts with a prompt that was served before, reuse that state and prefill only the remaining tokens, using prefill(suffix, past=...).
3. Memory budget. The cache may hold at most budget_tokens positions in total (the sum of kv.length over all entries). Evict least-recently-used entries to stay under the limit. A prompt longer than the budget is simply not cached.
4. Stats and isolation. Expose stats = {"hits", "misses", "tokens_saved"}. A hidden test sends the same prompt twice and expects identical outputs: make sure decoding a request can never corrupt what is stored in the cache.
5. Shared prefixes. Requests often share a long system prompt but end differently (SYS + question_a, SYS + question_b). Reuse the longest common prefix with any cached prompt, not just a cached prompt that is a full prefix. Lookup must be O(len(prompt)), not a scan over all entries. Add generate_batch(prompts, max_new_tokens) that returns outputs in input order.

Constraints

• Greedy decoding; outputs must equal the original Server token for token.
• 1 <= len(prompt) <= 4096, up to 10^4 requests, budget_tokens between 1 and 10^6.
• Tests compare model.prefill_tokens against an expected ceiling for each part.