Building a Cache-Augmented Generation System From Scratch

Most LLM tutorials open with a framework. You install a package, hand it a folder of documents, and twenty lines later you have a working chatbot. That's great if the goal is to ship something by Friday, but you come away with no idea what the model actually did with your text.

So I built a small Cache-Augmented Generation system the slow way. No LangChain, no LlamaIndex, no vector database, no hosted API. Just PyTorch, a 0.6B model and my own Ubuntu box with an RTX 4060 in it.

The path was roughly: check the machine, get GPU-enabled PyTorch installed, pull a small model down and run it locally, write a fake knowledge base, build an uncached baseline that works, then process the knowledge once and keep its KV cache around to answer many questions from.

By the end you have a real CAG prototype running on your own GPU. Small, but real.

What CAG actually is

Say you have a document that barely changes and a steady stream of people asking questions about it.

A normal grounded-generation setup does this:

Knowledge + Question 1 → Process everything → Answer

Knowledge + Question 2 → Process everything again → Answer

Knowledge + Question 3 → Process everything again → Answer

The knowledge is byte-for-byte identical each time, and the model recomputes all of it anyway.

CAG rearranges that:

Knowledge
    ↓
Process once
    ↓
Reusable KV cache
    ├── Question 1 → Answer
    ├── Question 2 → Answer
    └── Question 3 → Answer

A transformer's KV cache holds the attention keys and values for tokens it has already seen, so reusing them means you don't recompute them for every new token or every new request. Hugging Face calls this prefilling a cache, or prefix caching when the shared prompt gets reused across different generations.

Worth being clear about what this isn't. There's no training, no fine-tuning, and the weights never change.

1. Check the machine

I started in a fresh project directory:

mkdir -p ~/projects/protos/cag_proto
cd ~/projects/protos/cag_proto

Python version

python3 --version

Mine reported:

Python 3.13.7

Small thing that trips people up: don't write "python3 -- version" with a space. Python reads "version" as a filename it should go execute. The flag is one argument, "--version". And you don't need sudo for any of this.

NVIDIA GPU

nvidia-smi

That's the NVIDIA System Management Interface, and it prints the GPU model, driver version, memory usage, temperature and whatever processes are currently on the card.

Mine showed:

GPU:          NVIDIA GeForce RTX 4060
VRAM:         8188 MiB
Driver:       595.71.05
CUDA shown:   13.2

The desktop was sitting on about 430 MiB, so nearly all of the card was free.

RAM

free -h

free prints system memory and -h makes the numbers readable by humans instead of bytes.

Total RAM:       30 GiB
Available RAM:   27 GiB
Swap:             8 GiB

More than enough for a small local model.

2. Virtual environment

This keeps the project's libraries away from the system Python.

python3 -m venv .venv
source .venv/bin/activate

Your prompt should now start with (.venv). Confirm which Python you're actually running:

which python
python --version

The path should be inside the project:

/home/aatish/projects/protos/cag_proto/.venv/bin/python

Don't reach for sudo pip or sudo python in here. You'll end up with root-owned files and you'll quietly bypass the environment you just made.

3. Install PyTorch and the model libraries

Upgrade pip first:

python -m pip install --upgrade pip

The exact command that worked on this machine:

python -m pip install torch==2.12.1 \
  --index-url https://download.pytorch.org/whl/cu132

Don't copy that blindly onto a different machine. Go to PyTorch's install selector, pick your OS, package manager and compute platform, and use whatever it gives you. Their docs also suggest checking torch.cuda.is_available() afterwards, which we'll do in a second.

Then the rest:

python -m pip install --upgrade \
  numpy \
  transformers \
  accelerate \
  safetensors

What each one is for:

torch          Runs tensor operations and the model
numpy          General numerical operations
transformers   Loads the tokenizer and language model
accelerate     Helps place model components on available devices
safetensors    Loads model weights from the safetensors format

4. Confirm PyTorch can see the GPU

python -c "import torch; \
print('PyTorch:', torch.__version__); \
print('CUDA available:', torch.cuda.is_available()); \
print('GPU:', torch.cuda.get_device_name(0) \
if torch.cuda.is_available() else 'None')"

Mine printed:

PyTorch: 2.12.1+cu132
CUDA available: True
GPU: NVIDIA GeForce RTX 4060

"CUDA available: True" is the line that matters. Everything else is decoration.

Before I installed NumPy, PyTorch also threw this:

Failed to initialize NumPy: No module named 'numpy'

Which looks alarming next to CUDA output but has nothing to do with CUDA. The venv just didn't have NumPy in it yet.

5. Where the model comes from

The model identifier:

MODEL_NAME = "Qwen/Qwen3-0.6B"

That's not a folder on disk and it isn't a URL. It's a Hugging Face Hub repository ID, organization followed by repository name:

Organization: Qwen
Repository:   Qwen3-0.6B

When from_pretrained() gets that string it downloads the repository files and caches them locally, in a version-aware cache directory that looks like:

~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/

Poke around in it:

ls ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/

And check what it's costing you on disk:

du -sh ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/

The first run pulls down files like:

config.json
tokenizer_config.json
vocab.json
merges.txt
tokenizer.json
model.safetensors
generation_config.json

Every run after that reuses them.

You'll probably see this warning too:

You are sending unauthenticated requests to the HF Hub.
Please set a HF_TOKEN to enable higher rate limits and faster downloads.

Public models download fine without a token. You need one for higher rate limits or for gated repositories.

6. First local generation

Create test_model.py:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen3-0.6B"


def main() -> None:
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype="auto",
        device_map="auto",
    )

    print(f"Model device: {model.device}")

    messages = [
        {
            "role": "user",
            "content": "Explain cache-augmented generation in two simple sentences.",
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=False,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    input_length = inputs["input_ids"].shape[-1]
    print(f"Input tokens: {input_length}")

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    answer_tokens = generated_ids[0, input_length:]
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)

    print("\nModel answer:")
    print(answer)

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / (1024 ** 3)
        print(f"\nAllocated GPU memory: {allocated:.2f} GB")


if __name__ == "__main__":
    main()

Run it:

python test_model.py

The first run downloaded about 1.5 GB and gave me:

Model device: cuda:0
Input tokens: 24

Model answer:
Cache-augmented generation is a technique that improves
text generation by incorporating previously generated data...

Allocated GPU memory: 1.12 GB

That definition is wrong, which is a useful thing to hit on your very first run. Getting a model to load and generate says nothing about whether the output is correct. Qwen3-0.6B is tiny and I picked it to learn the plumbing, not because it knows what CAG is.

One flag worth explaining: Qwen3 has a thinking mode and a non-thinking mode, and enable_thinking=False in the chat template turns the thinking output off. Keeps the first prototype simple and a bit faster.

Where did the question come from?

It's hard-coded, right here:

messages = [
    {
        "role": "user",
        "content": "Explain cache-augmented generation in two simple sentences.",
    }
]

Nothing was typed into the terminal because this script isn't interactive yet. This line turns the messages into Qwen's chat format and then into token IDs:

inputs = tokenizer.apply_chat_template(...)

And this line produces the answer:

generated_ids = model.generate(...)

The whole round trip:

Human text
    ↓ tokenizer
Token IDs
    ↓ model
Generated token IDs
    ↓ tokenizer.decode()
Human-readable answer

7. A fictional knowledge base

I wanted facts the model could not possibly have picked up in training, so I invented some.

Create knowledge.txt:

Project Atish is an internal research project at Northstar Labs.

The project began on January 12, 2026.

Its objective is to build a solar-powered agricultural monitoring drone.

The project lead is Maya Sen.

The drone prototype is called Sparrow-7.

Sparrow-7 can fly for 47 minutes on a full charge.

Its main sensors measure soil moisture, air temperature, crop temperature,
and multispectral vegetation data.

The first field test took place in Mysuru.

During that test, Sparrow-7 covered 18 hectares and identified three irrigation
problems.

The next field test is scheduled to take place in Hubballi.

That file is the entire knowledge base.

8. The uncached baseline

Before caching anything you need something that works, so the baseline reads knowledge.txt, takes a question, jams both into one prompt, processes the lot and generates an answer.

Create baseline.py:

import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen3-0.6B"
KNOWLEDGE_FILE = Path("knowledge.txt")


def load_knowledge() -> str:
    """Read and validate the knowledge file."""
    if not KNOWLEDGE_FILE.exists():
        raise FileNotFoundError(f"Knowledge file not found: {KNOWLEDGE_FILE.resolve()}")

    knowledge = KNOWLEDGE_FILE.read_text(encoding="utf-8").strip()

    if not knowledge:
        raise ValueError("knowledge.txt is empty.")

    return knowledge


def main() -> None:
    print("Loading knowledge...")
    knowledge = load_knowledge()

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype="auto",
        device_map="auto",
    )
    model.eval()

    print(f"Model device: {model.device}")

    question = input("\nAsk a question about Project Atish: ").strip()

    if not question:
        print("No question was entered.")
        return

    messages = [
        {
            "role": "system",
            "content": (
                "Answer the user's question using only the supplied knowledge base. "
                "Give only the direct answer without introductory phrases. "
                "Do not use outside information. "
                "If the answer is not present, reply exactly: "
                "'The answer is not available in the knowledge base.'"
            ),
        },
        {
            "role": "user",
            "content": f"KNOWLEDGE BASE:\n{knowledge}\n\nQUESTION:\n{question}",
        },
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        enable_thinking=False,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    input_token_count = inputs["input_ids"].shape[-1]

    if torch.cuda.is_available():
        torch.cuda.synchronize()

    start_time = time.perf_counter()

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=120,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    if torch.cuda.is_available():
        torch.cuda.synchronize()

    elapsed_seconds = time.perf_counter() - start_time

    answer_tokens = generated_ids[0, input_token_count:]
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True).strip()

    print("\nAnswer:")
    print(answer)

    print("\nMeasurements:")
    print(f"Input tokens: {input_token_count}")
    print(f"Generated tokens: {answer_tokens.shape[-1]}")
    print(f"Generation time: {elapsed_seconds:.3f} seconds")

    if torch.cuda.is_available():
        allocated_gb = torch.cuda.memory_allocated() / (1024 ** 3)
        print(f"Allocated GPU memory: {allocated_gb:.2f} GB")


if __name__ == "__main__":
    main()

python baseline.py

Ask it something that's in the file:

What is Project Atish?

You should get roughly:

Project Atish is an internal research project at Northstar Labs.

Try another one:

How long can Sparrow-7 fly?

Sparrow-7 can fly for 47 minutes on a full charge.

Then ask for something that isn't in there at all:

What is the project's total budget?

The answer is not available in the knowledge base.

Both halves of that matter. The model answers from the document, and it declines when the document has nothing to say.

9. Why the baseline isn't CAG

It's grounded, but every single run still does the whole thing:

System instructions
+
Complete knowledge.txt
+
Question
        ↓
Tokenize everything
        ↓
Process everything
        ↓
Generate answer

Across three questions:

Question 1 → Process knowledge again
Question 2 → Process knowledge again
Question 3 → Process knowledge again

Nothing survives between questions. The script also exits after one question, so the model and all its runtime state vanish with the process.

Turning this into CAG means four changes. Split the reusable knowledge prefix away from the question that changes. Push the prefix through the model once. Keep the past_key_values that come back. Then give each question its own clean copy of that cache.

10. The CAG version

Hugging Face ships cache classes for this, DynamicCache and StaticCache. A DynamicCache grows as key and value states get appended, and you hand it to the model through the past_key_values parameter.

Create cag.py:

import copy
import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

MODEL_NAME = "Qwen/Qwen3-0.6B"
KNOWLEDGE_FILE = Path("knowledge.txt")
MAX_NEW_TOKENS = 120


def synchronize_gpu() -> None:
    """Wait for pending GPU operations."""
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def load_knowledge() -> str:
    """Read and validate knowledge.txt."""
    if not KNOWLEDGE_FILE.exists():
        raise FileNotFoundError(f"Could not find {KNOWLEDGE_FILE.resolve()}")

    knowledge = KNOWLEDGE_FILE.read_text(encoding="utf-8").strip()

    if not knowledge:
        raise ValueError("knowledge.txt is empty.")

    return knowledge


def main() -> None:
    knowledge = load_knowledge()

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype="auto",
        device_map="auto",
    )
    model.eval()

    print(f"Model device: {model.device}")

    system_message = {
        "role": "system",
        "content": (
            "Answer questions using only the supplied knowledge base. "
            "Give only the direct answer without introductory phrases. "
            "If the answer is missing, reply exactly: "
            "'The answer is not available in the knowledge base.'\n\n"
            f"KNOWLEDGE BASE:\n{knowledge}"
        ),
    }

    # Build the reusable knowledge prefix.
    prefix_text = tokenizer.apply_chat_template(
        [system_message],
        tokenize=False,
        add_generation_prompt=False,
        enable_thinking=False,
    )

    prefix_inputs = tokenizer(
        prefix_text,
        add_special_tokens=False,
        return_tensors="pt",
    ).to(model.device)

    prefix_token_count = prefix_inputs["input_ids"].shape[-1]

    print(f"Knowledge-prefix tokens: {prefix_token_count}")
    print("Building reusable KV cache...")

    # Begin with an empty cache.
    prefix_cache = DynamicCache(config=model.config)

    synchronize_gpu()
    cache_start = time.perf_counter()

    # Process the reusable knowledge once.
    with torch.inference_mode():
        cache_output = model(
            **prefix_inputs,
            past_key_values=prefix_cache,
            use_cache=True,
        )

    synchronize_gpu()
    cache_build_time = time.perf_counter() - cache_start

    # Save the populated cache.
    prefix_cache = cache_output.past_key_values

    print(f"Cache contains {prefix_cache.get_seq_length()} token positions.")
    print(f"Cache build time: {cache_build_time:.3f} seconds")

    print("\nCAG system ready.")
    print("Type 'exit' to stop.")

    while True:
        question = input("\nQuestion: ").strip()

        if question.lower() in {"exit", "quit"}:
            print("Exiting.")
            break

        if not question:
            print("Please enter a question.")
            continue

        messages = [system_message, {"role": "user", "content": question}]

        full_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )

        full_inputs = tokenizer(
            full_text,
            add_special_tokens=False,
            return_tensors="pt",
        ).to(model.device)

        full_input_length = full_inputs["input_ids"].shape[-1]

        # Confirm that the cached tokens are exactly the beginning of this prompt.
        cached_prefix_matches = torch.equal(
            full_inputs["input_ids"][:, :prefix_token_count],
            prefix_inputs["input_ids"],
        )

        if not cached_prefix_matches:
            raise RuntimeError("The cached token prefix does not match the question prompt.")

        # Generation extends the supplied cache. Give each question a clean copy.
        question_cache = copy.deepcopy(prefix_cache)

        synchronize_gpu()
        generation_start = time.perf_counter()

        with torch.inference_mode():
            generated_ids = model.generate(
                **full_inputs,
                past_key_values=question_cache,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                use_cache=True,
                pad_token_id=tokenizer.eos_token_id,
            )

        synchronize_gpu()
        generation_time = time.perf_counter() - generation_start

        answer_ids = generated_ids[0, full_input_length:]
        answer = tokenizer.decode(answer_ids, skip_special_tokens=True).strip()

        print("\nAnswer:")
        print(answer)

        question_token_count = full_input_length - prefix_token_count

        print("\nMeasurements:")
        print(f"Cached prefix tokens: {prefix_token_count}")
        print(f"Question-related tokens: {question_token_count}")
        print(f"Generated tokens: {answer_ids.shape[-1]}")
        print(f"Generation time: {generation_time:.3f} seconds")


if __name__ == "__main__":
    main()

python cag.py

Startup looks about like this:

Loading tokenizer...
Loading model...
Model device: cuda:0
Knowledge-prefix tokens: ...
Building reusable KV cache...
Cache contains ... token positions.
Cache build time: ... seconds

CAG system ready.
Type 'exit' to stop.

Now you can keep asking without restarting anything:

What is Project Atish?

How long can Sparrow-7 fly?

Where did the first field test take place?

What is the total budget?

Then exit when you're done.

11. The lines that make it CAG

Everything above is scaffolding. These five snippets are the actual idea.

Start with an empty cache:

prefix_cache = DynamicCache(config=model.config)

Run the knowledge prefix through the model once:

cache_output = model(
    **prefix_inputs,
    past_key_values=prefix_cache,
    use_cache=True,
)

Keep the keys and values that come back:

prefix_cache = cache_output.past_key_values

Hand each question a clean copy:

question_cache = copy.deepcopy(prefix_cache)

And generate against that copy:

generated_ids = model.generate(
    **full_inputs,
    past_key_values=question_cache,
    use_cache=True,
    ...
)

Hugging Face's own prefix-caching example does the same three moves, running a shared prompt through the model, copying the prefilled cache, and reusing it across different prompts.

So the shape of the thing is now:

knowledge.txt
     ↓
Tokenized once
     ↓
Model forward pass
     ↓
Reusable prefix KV cache
     ├── Copy + Question 1 → Answer
     ├── Copy + Question 2 → Answer
     └── Copy + Question 3 → Answer

12. Why the copy is not optional

Generation appends new key and value states to whatever cache you pass it. That's the part that bit me.

If question 1 gets the original cache, the cache afterwards holds:

Knowledge + Question 1 + Answer 1

Question 2 now starts from a polluted state, and every answer after that quietly builds on the last one. Hence the deepcopy per request. The original stays clean and each question gets its own branch:

Original cache:
Knowledge only

Question 1 cache:
Knowledge + Question 1 + Answer 1

Question 2 cache:
Knowledge + Question 2 + Answer 2

13. What's actually in the cache

Attention key and value tensors for the knowledge prefix. That's it.

What it is not:

• A vector database
• Search indexes
• Embedding vectors for retrieval
• Modified model weights
• Newly trained knowledge

The weights never move. The mental model I settled on:

Model weights
    = what the model learned during training

Knowledge document
    = external information provided at runtime

KV cache
    = the model's already-computed attention state
      for that runtime information

This is an inference-time trick and nothing more.

14. CAG against RAG

A typical RAG pipeline:

Question
    ↓
Search documents or vector store
    ↓
Select relevant chunks
    ↓
Insert chunks into prompt
    ↓
Generate answer

The prototype above:

Load complete manageable knowledge base
    ↓
Build its KV cache
    ↓
Reuse that cache for every question

CAG takes retrieval out of the request path entirely. That's a good trade when the knowledge is stable, when the whole thing fits in the model's usable context, when lots of questions share the same knowledge, and when you care about latency on the repeated prefix.

It's a bad trade when the knowledge changes constantly, when the corpus is far bigger than the context window, when different users need completely different knowledge, or when you need document-level citations pointing at exactly where an answer came from.

Neither one replaces the other. They solve different operational problems and there's nothing stopping you from combining them.

15. Why this example doesn't look faster

My knowledge file is about ten lines long. Prefilling that was never expensive, so the gap between the baseline and the CAG version is small enough to disappear into measurement noise.

The architectural difference only shows up once you have hundreds or thousands of tokens of prefix, repeated questions, a warmed-up model and several benchmark runs to average over.

A real benchmark would separate out:

• Cache construction time
• Question-prefill time
• First-token latency
• Total generation time
• Generated tokens per second
• GPU memory used by the cache

And it would run each scenario multiple times rather than trusting one number.

16. What this prototype can't do

It's genuine CAG, and it's nowhere near production.

The cache lives in GPU memory only, so when cag.py exits:

Process ends
    ↓
GPU memory is released
    ↓
KV cache disappears

Start it again and you rebuild from scratch.

The rest of the list: knowledge.txt has to fit inside the context window, any change to the knowledge means rebuilding the cache, questions are handled one at a time, there's no HTTP API, no web interface, no automated evaluation of the answers, no persistence across restarts and no cache isolation between users. On top of that, a 0.6B model will misread anything complicated, and there's no citation or evidence extraction anywhere in the loop.

17. One small terminal thing

Running the model in one terminal while you watch nvidia-smi in another is genuinely useful. On the terminal I was using, Ctrl+Shift+N opened a separate window and Ctrl+Shift+T opened another tab. Varies by emulator, so check yours.

18. Where I'm taking it next

Roughly in this order:

Current command-line CAG prototype
        ↓
Add repeated benchmarking
        ↓
Load larger text and Markdown files
        ↓
Calculate cache memory usage
        ↓
Add cache invalidation when files change
        ↓
Expose the system through FastAPI
        ↓
Build a small chat interface
        ↓
Add concurrency and per-user sessions
        ↓
Compare CAG against RAG

I'm going to resist pulling in frameworks for as long as I can. The whole point of the first version was to see the chain end to end:

Text
    ↓
Tokens
    ↓
Model forward pass
    ↓
Attention keys and values
    ↓
Reusable prefix cache
    ↓
Question-specific generation

Once you've watched that happen, evaluating the higher-level libraries gets a lot easier, because you know what they're hiding.

Wrapping up

Started with nothing installed and ended with a working pipeline: Ubuntu box, virtual environment, GPU-enabled PyTorch, Qwen3-0.6B pulled down locally, a first generation, a fake knowledge base, a grounded baseline, a reusable KV cache and an interactive prototype sitting on top of it.

The conceptual move was going from this:

Knowledge + question processed repeatedly

to this:

Knowledge processed once
+
Question-specific generation

Small system. Real one.