AKV

Cache Engineering for Agents

🌐 Project Page • 📑 Paper

AKV brings KV cache engineering to LLM agents. It combines agent-side policies for deciding what to drop and when to compact positions with an SGLang-based inference system that executes those operations while reusing compatible KV cache.

• akv.agent is the algorithm component that decides which messages to drop from the KV cache and when to reposition retained tokens.
• akv.system executes cache operations through SGLang's OpenAI-compatible Chat Completions API.

If you find AKV useful, please consider citing our paper:

@article{fu2026akv,
    title={Cache Engineering for Agents},
    author={Tianyu Fu and Ruoxi Wang and Zihan Min and Zhongqi Ren and Zhihao Li and Tengxuan Liu and Shengxu Mu and Hanling Zhang and Yu Wang},
    journal={arXiv preprint arXiv:2610.28577},
    year={2026},
}

curl -o install-rust.sh https://sh.rustup.rs
sh install-rust.sh -y
source "$HOME/.cargo/env"

git clone https://github.com/thu-nics/AKV.git
cd AKV
uv venv --python 3.12
source .venv/bin/activate
uv pip install -e .

The environment is stored in .venv to avoid conflicting with the akv/ source directory. See the system README and the agent README for dependencies.

These examples use GPT-OSS-120B, served as gpt-oss-120b.

python -m akv.system.launch_server \
  --model-path /path/to/gpt-oss-120b \
  --reasoning-parser gpt-oss \
  --tool-call-parser gpt-oss \
  --tp-size 2 \
  --port 8000

For additional launch options, see the system README.

AKV uses CachePolicy to determine cache operations based on the current request and execution history. Given an OpenAI-compatible request payload, CachePolicy augments it with cache operation fields as needed: drop_message specifies which messages' KV entries to drop and when to apply Drop, while reposition specifies when to apply Reposition to compact the retained token positions.

The following example retains KV entries for tool results from the latest completed turn and schedules Reposition. Here, a turn consists of one assistant response that calls tools and all of its tool results.

from akv.agent import CachePolicy
from akv.agent.policies import TurnTrigger, TurnSelection, TurnReposition

payload = {
    "model": "gpt-oss-120b",
    "messages": [
        {
            "role": "user", 
            "content": "Compare records A and B."
        },
        {
            "role": "assistant", 
            "tool_calls": [{
                "id": "a", "type": "function",
                "function": {"name": "read_record", "arguments": '{"record_id":"A"}'},
            }],
        },
        {
            "role": "tool", 
            "tool_call_id": "a", 
            "content": "Record A"
        },
        {
            "role": "assistant", 
            "tool_calls": [{
                "id": "b", "type": "function",
                "function": {"name": "read_record", "arguments": '{"record_id":"B"}'},
            }],
        },
        {
            "role": "tool", 
            "tool_call_id": "b", 
            "content": "Record B"
        },
    ],
    "stream": True,
    "max_completion_tokens": 2048,
}

policy = CachePolicy(
    trigger=TurnTrigger(),
    selector=TurnSelection(1),
    reposition=TurnReposition(after=2, every=2),
)
akv_payload, cache_operation_history = policy.process(payload, api_style="sglang")

See Customizing CachePolicy for details on the three components: trigger, selector, and reposition.

The request contains two completed tool calls, with results at message indices 2 and 4. akv_payload preserves the original payload and adds the following cache operations:

{ "drop_message": {"4": [2]}, "reposition": [4] }

These operations instruct AKV System to drop the earlier tool result (message 2) after processing the new result (message 4), then compact the remaining token positions.

akv_payload = {
    "model": "gpt-oss-120b",
    "messages": [
        {
            "role": "user",
            "content": "Compare records A and B."
        },
        {
            "role": "assistant",
            "tool_calls": [{
                "id": "a", "type": "function",
                "function": {"name": "read_record", "arguments": '{"record_id":"A"}'},
            }],
        },
        {
            "role": "tool",
            "tool_call_id": "a",
            "content": "Record A"
        },
        {
            "role": "assistant",
            "tool_calls": [{
                "id": "b", "type": "function",
                "function": {"name": "read_record", "arguments": '{"record_id":"B"}'},
            }],
        },
        {
            "role": "tool",
            "tool_call_id": "b",
            "content": "Record B"
        },
    ],
    "drop_message": {"4": [2]},
    "reposition": [4],
    "stream": True,
    "max_completion_tokens": 2048,
}

After applying the policy, you can send akv_payload directly to the system as the request body.

import requests

response = requests.post(
    "http://localhost:8000/v1/chat/completions",
    json=akv_payload,
    timeout=120,
)
response.raise_for_status()

After processing message 4, AKV System removes message 2's KV entries, compacts the retained token positions, and generates the next reply using the remaining context.

The demo provides a complete example of an agent making tool calls over multiple turns. The agent follows a chain of delivery records to find a parcel’s final delivery location.

AKV’s CacheTool interface lets agents manage the KV cache through tool calls. In this demo, DropContextTool exposes drop_context, which the agent can call to schedule Drop for tool results it no longer needs. Before the next model call, CachePolicy converts successful drop_context calls into drop_message operations. Reposition is scheduled after every three completed tool-calling turns.

Run the demo with:

python examples/demo.py \
  --base-url http://localhost:8000/v1 \
  --model gpt-oss-120b \
  --output local/demo-tool-01

As the demo runs, you can watch the model follow the chain of delivery records, manage its context by discarding results it no longer needs, and ultimately identify the parcel's final delivery location.

Integrating AKV requires only one step: process the final payload before sending request to the system. After the harness serializes its history and tools, pass the payload to CachePolicy.process, then send the returned payload to AKV System.

Optional: register cache tools. If your policy uses a CacheTool, register its schema and execute handler with the harness before starting the agent loop.

CachePolicy manages the KV cache throughout agent execution by scheduling Drop and Reposition. Drop removes selected messages' KV entries without changing the retained token positions, while Reposition compacts those positions. Retained KV entries are reused without re-prefilling after both operations.

policy = CachePolicy(
    trigger=TurnTrigger(),
    selector=TurnSelection(2),
    reposition=TurnReposition(after=4, every=2),
)

CachePolicy separates when to drop, what to drop, and whether to reposition. Each component receives a PolicyContext containing the current request and cache operation history.

The built-in components support three approaches:

• Turn-based: trigger Drop and Reposition at specified turn intervals, and retain context from the latest turns.
• Token-based: use token counts to decide when to drop or reposition, and how much context to retain. These components require a token counter matching the model.
• Model-based: let the model decide what to drop and whether to reposition through a CacheTool.

To customize a policy, supply your own trigger, selector, or reposition function to CachePolicy.

For example, this policy checks after each completed tool-calling turn, selects all but the latest visible tool result for Drop, and schedules Reposition whenever the selection is nonempty:

from akv.agent import CachePolicy, PolicyContext


def after_tool_turn(context: PolicyContext) -> bool:
    return context.new_turn


def select_old_results(context: PolicyContext) -> set[int]:
    visible_results = [
        i for i, message in enumerate(context.request.messages)
        if message["role"] == "tool" and i not in context.hidden
    ]
    return set(visible_results[:-1])  # Keep the latest result.


def reposition_after_drop(context: PolicyContext, selected: set[int]) -> bool:
    return bool(selected)


policy = CachePolicy(
    trigger=after_tool_turn,
    selector=select_old_results,
    reposition=reposition_after_drop,
)

With the Quick Start payload, this policy selects message 2 for Drop, retains message 4's KV entries, and schedules Reposition. It adds drop_message={"4": [2]} and reposition=[4] to the request.

Context requests report cached_tokens, drop_skipped_tokens, and repos_tokens in usage.prompt_tokens_details, including zero values. See the system README for rule details and usage accounting.

Drop/Reposition is currently enabled for these text-model architectures:

See the system README for runtime requirements and compatibility details.

Explore more efficient LLM projects from us:

R2R Token-level routing for reasoning LLMs

TkR Efficient serving for token-level LLM routing

TaH Adaptive Looped LLMs

C2C Communicate through KV-Cache between LLMs

FrF Efficient video token reduction for LVLMs

MoA Mixture of sparse attention for LLMs