Building an Attention-Driven Semantic Caching Layer for LLM APIs

TL;DR — Build a semantic caching layer for LLM APIs that uses self-attention mechanisms to score query relevance, drastically reducing latency and compute costs. This project demonstrates advanced systems skills in memory management, concurrency, and vector indexing that hiring managers actively seek.

The bottleneck in modern LLM applications isn’t the model’s intelligence; it’s the inference latency and cost. Standard caching mechanisms like Least Recently Used (LRU) fail catastrophically in semantic spaces because a query for “how to bake a cake” is fundamentally different from “recipe instructions,” even if they share similar recency. To solve this, we need a cache that understands meaning.

Enter the Attention-Driven Semantic Cache. By leveraging self-attention mechanisms, we can dynamically score the relevance of incoming queries against a history of prompts and responses. This project is a hands-on build guide for constructing this infrastructure, proving you can bridge the gap between machine learning theory and distributed systems engineering.

Why This Project Stands Out on a CV

Hiring managers for ML Infrastructure and Backend roles are constantly searching for candidates who understand the intersection of high-throughput systems and modern AI. This project signals three critical competencies:

Distributed Systems & Concurrency: Implementing an async, thread-safe cache requires a deep understanding of race conditions, locks, and event loops (e.g., Python’s asyncio).
Memory Management: Building a paged memory buffer demonstrates you know how to manage finite hardware resources efficiently, a core tenet of systems programming.
ML Infrastructure: Moving beyond simply calling an API to optimizing the routing and retrieval of LLM outputs shows you can build scalable AI products, not just prototypes.

This project positions you for roles like ML Infrastructure Engineer, Backend SRE, or Applied Systems Engineer, where the focus is on making models performant and cost-effective in production.

Architecture Overview

The system is designed as a decoupled, high-throughput pipeline. It separates the concerns of embedding generation, attention scoring, and cache storage. Here is how the components fit together:

API Gateway: Receives incoming LLM prompts and intercepts requests before they hit the heavy inference model.
Embedding Engine: Transforms the raw text prompt into a dense vector representation using a pre-trained model like sentence-transformers.
Attention Scorer: Computes the self-attention relevance score between the incoming query vector and the stored keys in the cache. This determines if a cached response is semantically similar enough to reuse.
Paged Memory Manager: Stores the vector keys and LLM responses in a paged buffer, mimicking the virtual memory management found in modern OS kernels to prevent memory fragmentation.
Eviction Policy: A hybrid policy that removes least-relevant pages based on attention scores when the memory limit is reached.

   [API Gateway] 
         |
         v
   [Embedding Engine] ---> [Attention Scorer] 
         |                      |
         v                      v
   [Paged Memory Manager] <-- [Eviction Policy]
         |
         v
   [Cached Response / LLM Inference]

Building It Step by Step

We will build this using Python, leveraging torch for the attention mechanics, sentence-transformers for embeddings, and asyncio for concurrency. The core logic relies on calculating the dot-product attention between the incoming query and the cached keys.

Step 1: Define the Paged Memory Structure

First, we need a memory manager that allocates fixed-size pages. This prevents memory fragmentation and allows for O(1) insertion and deletion, similar to the paged attention mechanism used in vLLM.

import asyncio
from collections import OrderedDict
import numpy as np
from typing import Optional, Tuple

class PagedMemoryManager:
    def __init__(self, max_pages: int, page_size: int, embedding_dim: int):
        self.max_pages = max_pages
        self.page_size = page_size
        self.embedding_dim = embedding_dim
        self.pages = OrderedDict() # Acts as our physical memory block
        self.lock = asyncio.Lock()

    async def allocate(self, key: np.ndarray, value: str) -> bool:
        async with self.lock:
            if len(self.pages) >= self.max_pages:
                return False
            self.pages[key.tobytes()] = {
                "key": key,
                "value": value,
                "access_time": asyncio.get_event_loop().time()
            }
            return True

    def get_all_keys(self) -> Tuple[np.ndarray, list]:
        keys = np.array([v["key"] for v in self.pages.values()])
        values = [v["value"] for v in self.pages.values()]
        return keys, values

Step 2: Implement the Attention Scoring Mechanism

Next, we implement the core self-attention logic. Instead of the full softmax attention used in transformers, we use a simplified scaled dot-product attention to score the relevance of the incoming query against our cached keys.

import torch
import torch.nn.functional as F

class AttentionScorer:
    def __init__(self, scaling_factor: float = 0.1):
        self.scaling_factor = scaling_factor

    def calculate_scores(self, query: np.ndarray, cached_keys: np.ndarray) -> np.ndarray:
        # Convert to torch tensors for GPU/CPU acceleration
        q_tensor = torch.tensor(query, dtype=torch.float32)
        k_tensor = torch.tensor(cached_keys, dtype=torch.float32)

        # Scaled Dot-Product Attention: softmax(Q @ K.T / sqrt(d_k))
        scores = torch.matmul(q_tensor, k_tensor.T) * self.scaling_factor
        probabilities = F.softmax(scores, dim=-1)
        
        return probabilities.detach().numpy()

Step 3: Integrate the Semantic Cache

Now we combine the memory manager and the scorer into a unified cache. If the attention score for a cached key exceeds a threshold, we return the cached value. Otherwise, we route the request to the LLM and store the result.

class AttentionSemanticCache:
    def __init__(self, memory_manager: PagedMemoryManager, scorer: AttentionScorer, threshold: float = 0.85):
        self.memory_manager = memory_manager
        self.scorer = scorer
        self.threshold = threshold

    async def retrieve(self, query_embedding: np.ndarray) -> Optional[str]:
        keys, values = self.memory_manager.get_all_keys()
        if len(keys) == 0:
            return None

        scores = self.scorer.calculate_scores(query_embedding, keys)
        max_score = np.max(scores)
        best_idx = np.argmax(scores)

        if max_score >= self.threshold:
            # Update access time for LRU eviction
            best_key = keys[best_idx].tobytes()
            self.memory_manager.pages[best_key]["access_time"] = asyncio.get_event_loop().time()
            return values[best_idx]
        
        return None

    async def store(self, query_embedding: np.ndarray, response: str):
        await self.memory_manager.allocate(query_embedding, response)

Running and Testing It

To prove this system works under load, we will set up a local test suite using pytest and a load testing tool like locust. This ensures our async locks are functioning correctly and the attention scoring is returning the expected semantic matches.

First, initialize your environment and install dependencies:

pip install torch sentence-transformers numpy pytest locust

Next, write a test to verify the cache hit rate and attention threshold logic:

import pytest
import numpy as np
from your_cache_module import PagedMemoryManager, AttentionScorer, AttentionSemanticCache

@pytest.mark.asyncio
async def test_attention_cache_hit():
    manager = PagedMemoryManager(max_pages=10, page_size=512, embedding_dim=384)
    scorer = AttentionScorer(scaling_factor=0.1)
    cache = AttentionSemanticCache(manager, scorer, threshold=0.5)

    # Store a known embedding and response
    query = np.random.rand(384)
    await cache.store(query, "Cached Response")

    # Retrieve with a highly similar query
    similar_query = query + np.random.normal(0, 0.01, 384)
    result = await cache.retrieve(similar_query)

    assert result == "Cached Response"

To run the load test, create a locustfile.py that sends concurrent requests to your cache:

from locust import HttpUser, task, between
import numpy as np

class CacheUser(HttpUser):
    wait_time = between(0.1, 0.5)

    @task
    def query_cache(self):
        query = np.random.rand(384).tolist()
        self.client.post("/cache/retrieve", json={"embedding": query})

Run the test with locust -f locustfile.py and navigate to http://localhost:8089 to simulate hundreds of concurrent users, verifying that the async lock prevents race conditions.

Extending It: Your Roadmap to Senior-Level

Moving this project from a toy example to a production-grade system requires implementing features that hiring managers associate with senior engineers. Here are six concrete upgrades:

Persistence via RocksDB: Integrate RocksDB to persist the paged memory to disk. Why it matters: It ensures fault tolerance and crash recovery, preventing total cache loss on server restart.
Horizontal Scaling with gRPC: Expose the cache as a microservice using gRPC and share state via Redis. Why it matters: It allows the system to handle increased throughput by distributing the cache across multiple nodes.
Observability with OpenTelemetry: Add OpenTelemetry metrics to track cache hit rates, attention score distributions, and latency. Why it matters: It provides the visibility required to debug performance bottlenecks in a production environment.
Quantized Key Storage: Store the embedding keys in INT8 or FP8 formats using libraries like bitsandbytes. Why it matters: It drastically reduces memory footprint, allowing the cache to hold significantly more vectors in the same RAM.
Continuous Batching: Implement a batching mechanism that groups incoming queries before calculating attention scores. Why it matters: It maximizes GPU utilization during the scoring phase, lowering the per-query latency.
Hybrid LRU-Attention Eviction: Combine the attention score with an LRU timestamp to decide which pages to evict. Why it matters: It prevents the cache from being dominated by highly relevant but stale data, ensuring robust long-term performance.

Key Takeaways

Standard caching fails in semantic spaces; self-attention provides a mathematically sound method for relevance scoring.
Paged memory management is critical for preventing fragmentation and ensuring O(1) operations in high-throughput systems.
Async concurrency and locking mechanisms are mandatory when sharing state between LLM inference and cache retrieval.
Bridging ML concepts (attention) with systems concepts (paging, gRPC) is the defining skill gap for modern ML Infrastructure Engineers.
Production readiness requires moving beyond in-memory storage to persistent, observable, and horizontally scalable architectures.

Why This Project Stands Out on a CV#

Architecture Overview#

Building It Step by Step#

Step 1: Define the Paged Memory Structure#

Step 2: Implement the Attention Scoring Mechanism#

Step 3: Integrate the Semantic Cache#

Running and Testing It#

Extending It: Your Roadmap to Senior-Level#

Key Takeaways#

Further Reading#