KV Caching Explained: A Deep Dive into Optimizing Transformer Inference

Introduction to KV Caching

When large language models (LLMs) generate text autoregressively, they perform redundant computations by reprocessing the same tokens repeatedly. Key-Value (KV) Caching solves this by storing intermediate attention states, dramatically improving inference speed – often by 5x or more in practice.

In this comprehensive guide, we’ll:

Explain the transformer attention bottleneck
Implement KV caching from scratch in PyTorch
Benchmark performance gains
Compare with Hugging Face’s built-in implementation
Discuss advanced optimizations like grouped-query attention

1. The Transformer Attention Bottleneck

Standard Autoregressive Inference

Without caching, each new token generation requires:

<span class="token comment"># Pseudocode: Naive generation</span>
<span class="token keyword">for</span> token <span class="token keyword">in</span> output_sequence<span class="token punctuation">:</span>
    <span class="token comment"># Reprocess ENTIRE sequence each time!</span>
    output <span class="token operator">=</span> model<span class="token punctuation">(</span>input_sequence <span class="token operator">+</span> generated_tokens<span class="token punctuation">)</span> 
    next_token <span class="token operator">=</span> sample<span class="token punctuation">(</span>output<span class="token punctuation">[</span><span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">]</span><span class="token punctuation">)</span>

# Pseudocode: Naive generation

for token in output_sequence:

# Reprocess ENTIRE sequence each time!

output = model(input_sequence + generated_tokens)

next_token = sample(output[-1])

Problem: For sequence length N, this results in O(N²) computations due to:

Repeated matrix multiplications for Q/K/V
Full attention score recalculations

Attention Mechanics Refresher

Each transformer layer computes:

Q <span class="token operator">=</span> X @ W_q  <span class="token comment"># (seq_len, d_head)</span>
K <span class="token operator">=</span> X @ W_k  <span class="token comment"># (seq_len, d_head)</span>
V <span class="token operator">=</span> X @ W_v  <span class="token comment"># (seq_len, d_head)</span>

attn <span class="token operator">=</span> softmax<span class="token punctuation">(</span>Q @ K<span class="token punctuation">.</span>T <span class="token operator">/</span> sqrt<span class="token punctuation">(</span>d_head<span class="token punctuation">)</span><span class="token punctuation">)</span> @ V

Q = X @ W_q # (seq_len, d_head)

K = X @ W_k # (seq_len, d_head)

V = X @ W_v # (seq_len, d_head)

attn = softmax(Q @ K.T / sqrt(d_head)) @ V

2. Implementing KV Cache from Scratch

Complete PyTorch Implementation

<span class="token keyword">import</span> torch
<span class="token keyword">import</span> torch<span class="token punctuation">.</span>nn <span class="token keyword">as</span> nn
<span class="token keyword">from</span> typing <span class="token keyword">import</span> Dict<span class="token punctuation">,</span> Optional

<span class="token keyword">class</span> <span class="token class-name">KVCache</span><span class="token punctuation">:</span>
    <span class="token keyword">def</span> <span class="token function">__init__</span><span class="token punctuation">(</span>self<span class="token punctuation">,</span> max_batch_size<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">,</span> max_seq_len<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">,</span> 
                 n_heads<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">,</span> head_dim<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">,</span> device<span class="token punctuation">:</span> <span class="token builtin">str</span> <span class="token operator">=</span> <span class="token string">"cuda"</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
        self<span class="token punctuation">.</span>cache <span class="token operator">=</span> <span class="token punctuation">{</span>
            <span class="token string">"key"</span><span class="token punctuation">:</span> torch<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span>
                <span class="token punctuation">(</span>max_batch_size<span class="token punctuation">,</span> n_heads<span class="token punctuation">,</span> max_seq_len<span class="token punctuation">,</span> head_dim<span class="token punctuation">)</span><span class="token punctuation">,</span>
                device<span class="token operator">=</span>device<span class="token punctuation">)</span><span class="token punctuation">,</span>
            <span class="token string">"value"</span><span class="token punctuation">:</span> torch<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span>
                <span class="token punctuation">(</span>max_batch_size<span class="token punctuation">,</span> n_heads<span class="token punctuation">,</span> max_seq_len<span class="token punctuation">,</span> head_dim<span class="token punctuation">)</span><span class="token punctuation">,</span>
                device<span class="token operator">=</span>device<span class="token punctuation">)</span>
        <span class="token punctuation">}</span>
        self<span class="token punctuation">.</span>position <span class="token operator">=</span> <span class="token number">0</span>
    
    <span class="token keyword">def</span> <span class="token function">update</span><span class="token punctuation">(</span>self<span class="token punctuation">,</span> 
               new_key<span class="token punctuation">:</span> torch<span class="token punctuation">.</span>Tensor<span class="token punctuation">,</span>  <span class="token comment"># (batch, n_heads, new_tokens, head_dim)</span>
               new_value<span class="token punctuation">:</span> torch<span class="token punctuation">.</span>Tensor<span class="token punctuation">)</span> <span class="token operator">-</span><span class="token operator">></span> Dict<span class="token punctuation">[</span><span class="token builtin">str</span><span class="token punctuation">,</span> torch<span class="token punctuation">.</span>Tensor<span class="token punctuation">]</span><span class="token punctuation">:</span>
        
        batch_size <span class="token operator">=</span> new_key<span class="token punctuation">.</span>size<span class="token punctuation">(</span><span class="token number">0</span><span class="token punctuation">)</span>
        seq_len <span class="token operator">=</span> new_key<span class="token punctuation">.</span>size<span class="token punctuation">(</span><span class="token number">2</span><span class="token punctuation">)</span>
        
        <span class="token comment"># Update cache slices</span>
        self<span class="token punctuation">.</span>cache<span class="token punctuation">[</span><span class="token string">"key"</span><span class="token punctuation">]</span><span class="token punctuation">[</span><span class="token punctuation">:</span>batch_size<span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> self<span class="token punctuation">.</span>position<span class="token punctuation">:</span>self<span class="token punctuation">.</span>position<span class="token operator">+</span>seq_len<span class="token punctuation">]</span> <span class="token operator">=</span> new_key
        self<span class="token punctuation">.</span>cache<span class="token punctuation">[</span><span class="token string">"value"</span><span class="token punctuation">]</span><span class="token punctuation">[</span><span class="token punctuation">:</span>batch_size<span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> self<span class="token punctuation">.</span>position<span class="token punctuation">:</span>self<span class="token punctuation">.</span>position<span class="token operator">+</span>seq_len<span class="token punctuation">]</span> <span class="token operator">=</span> new_value
        
        self<span class="token punctuation">.</span>position <span class="token operator">+=</span> seq_len
        
        <span class="token keyword">return</span> <span class="token punctuation">{</span>
            <span class="token string">"key"</span><span class="token punctuation">:</span> self<span class="token punctuation">.</span>cache<span class="token punctuation">[</span><span class="token string">"key"</span><span class="token punctuation">]</span><span class="token punctuation">[</span><span class="token punctuation">:</span>batch_size<span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> <span class="token punctuation">:</span>self<span class="token punctuation">.</span>position<span class="token punctuation">]</span><span class="token punctuation">,</span>
            <span class="token string">"value"</span><span class="token punctuation">:</span> self<span class="token punctuation">.</span>cache<span class="token punctuation">[</span><span class="token string">"value"</span><span class="token punctuation">]</span><span class="token punctuation">[</span><span class="token punctuation">:</span>batch_size<span class="token punctuation">,</span> <span class="token punctuation">:</span><span class="token punctuation">,</span> <span class="token punctuation">:</span>self<span class="token punctuation">.</span>position<span class="token punctuation">]</span>
        <span class="token punctuation">}</span>

<span class="token keyword">class</span> <span class="token class-name">CausalSelfAttention</span><span class="token punctuation">(</span>nn<span class="token punctuation">.</span>Module<span class="token punctuation">)</span><span class="token punctuation">:</span>
    <span class="token keyword">def</span> <span class="token function">__init__</span><span class="token punctuation">(</span>self<span class="token punctuation">,</span> d_model<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">,</span> n_heads<span class="token punctuation">:</span> <span class="token builtin">int</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
        <span class="token builtin">super</span><span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span>__init__<span class="token punctuation">(</span><span class="token punctuation">)</span>
        self<span class="token punctuation">.</span>d_head <span class="token operator">=</span> d_model <span class="token operator">//</span> n_heads
        self<span class="token punctuation">.</span>n_heads <span class="token operator">=</span> n_heads
        self<span class="token punctuation">.</span>W_q <span class="token operator">=</span> nn<span class="token punctuation">.</span>Linear<span class="token punctuation">(</span>d_model<span class="token punctuation">,</span> d_model<span class="token punctuation">)</span>
        self<span class="token punctuation">.</span>W_k <span class="token operator">=</span> nn<span class="token punctuation">.</span>Linear<span class="token punctuation">(</span>d_model<span class="token punctuation">,</span> d_model<span class="token punctuation">)</span>
        self<span class="token punctuation">.</span>W_v <span class="token operator">=</span> nn<span class="token punctuation">.</span>Linear<span class="token punctuation">(</span>d_model<span class="token punctuation">,</span> d_model<span class="token punctuation">)</span>
        
    <span class="token keyword">def</span> <span class="token function">forward</span><span class="token punctuation">(</span>self<span class="token punctuation">,</span> 
                x<span class="token punctuation">:</span> torch<span class="token punctuation">.</span>Tensor<span class="token punctuation">,</span>
                kv_cache<span class="token punctuation">:</span> Optional<span class="token punctuation">[</span>KVCache<span class="token punctuation">]</span> <span class="token operator">=</span> <span class="token boolean">None</span><span class="token punctuation">)</span> <span class="token operator">-</span><span class="token operator">></span> torch<span class="token punctuation">.</span>Tensor<span class="token punctuation">:</span>
        
        batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> _ <span class="token operator">=</span> x<span class="token punctuation">.</span>shape
        
        <span class="token comment"># Project to Q/K/V</span>
        Q <span class="token operator">=</span> self<span class="token punctuation">.</span>W_q<span class="token punctuation">(</span>x<span class="token punctuation">)</span><span class="token punctuation">.</span>view<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> self<span class="token punctuation">.</span>n_heads<span class="token punctuation">,</span> self<span class="token punctuation">.</span>d_head<span class="token punctuation">)</span><span class="token punctuation">.</span>transpose<span class="token punctuation">(</span><span class="token number">1</span><span class="token punctuation">,</span> <span class="token number">2</span><span class="token punctuation">)</span>
        K <span class="token operator">=</span> self<span class="token punctuation">.</span>W_k<span class="token punctuation">(</span>x<span class="token punctuation">)</span><span class="token punctuation">.</span>view<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> self<span class="token punctuation">.</span>n_heads<span class="token punctuation">,</span> self<span class="token punctuation">.</span>d_head<span class="token punctuation">)</span><span class="token punctuation">.</span>transpose<span class="token punctuation">(</span><span class="token number">1</span><span class="token punctuation">,</span> <span class="token number">2</span><span class="token punctuation">)</span>
        V <span class="token operator">=</span> self<span class="token punctuation">.</span>W_v<span class="token punctuation">(</span>x<span class="token punctuation">)</span><span class="token punctuation">.</span>view<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> self<span class="token punctuation">.</span>n_heads<span class="token punctuation">,</span> self<span class="token punctuation">.</span>d_head<span class="token punctuation">)</span><span class="token punctuation">.</span>transpose<span class="token punctuation">(</span><span class="token number">1</span><span class="token punctuation">,</span> <span class="token number">2</span><span class="token punctuation">)</span>
        
        <span class="token comment"># Update cache if provided</span>
        <span class="token keyword">if</span> kv_cache <span class="token keyword">is</span> <span class="token keyword">not</span> <span class="token boolean">None</span><span class="token punctuation">:</span>
            cache <span class="token operator">=</span> kv_cache<span class="token punctuation">.</span>update<span class="token punctuation">(</span>K<span class="token punctuation">,</span> V<span class="token punctuation">)</span>
            K<span class="token punctuation">,</span> V <span class="token operator">=</span> cache<span class="token punctuation">[</span><span class="token string">"key"</span><span class="token punctuation">]</span><span class="token punctuation">,</span> cache<span class="token punctuation">[</span><span class="token string">"value"</span><span class="token punctuation">]</span>
        
        <span class="token comment"># Scaled dot-product attention</span>
        attn <span class="token operator">=</span> <span class="token punctuation">(</span>Q @ K<span class="token punctuation">.</span>transpose<span class="token punctuation">(</span><span class="token operator">-</span><span class="token number">2</span><span class="token punctuation">,</span> <span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">)</span><span class="token punctuation">)</span> <span class="token operator">/</span> torch<span class="token punctuation">.</span>sqrt<span class="token punctuation">(</span>torch<span class="token punctuation">.</span>tensor<span class="token punctuation">(</span>self<span class="token punctuation">.</span>d_head<span class="token punctuation">)</span><span class="token punctuation">)</span>
        mask <span class="token operator">=</span> torch<span class="token punctuation">.</span>tril<span class="token punctuation">(</span>torch<span class="token punctuation">.</span>ones<span class="token punctuation">(</span>seq_len<span class="token punctuation">,</span> seq_len<span class="token punctuation">)</span><span class="token punctuation">)</span><span class="token punctuation">.</span>to<span class="token punctuation">(</span>x<span class="token punctuation">.</span>device<span class="token punctuation">)</span>
        attn <span class="token operator">=</span> attn<span class="token punctuation">.</span>masked_fill<span class="token punctuation">(</span>mask <span class="token operator">==</span> <span class="token number">0</span><span class="token punctuation">,</span> <span class="token builtin">float</span><span class="token punctuation">(</span><span class="token string">'-inf'</span><span class="token punctuation">)</span><span class="token punctuation">)</span>
        attn <span class="token operator">=</span> torch<span class="token punctuation">.</span>softmax<span class="token punctuation">(</span>attn<span class="token punctuation">,</span> dim<span class="token operator">=</span><span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">)</span> @ V
        
        <span class="token keyword">return</span> attn<span class="token punctuation">.</span>transpose<span class="token punctuation">(</span><span class="token number">1</span><span class="token punctuation">,</span> <span class="token number">2</span><span class="token punctuation">)</span><span class="token punctuation">.</span>reshape<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> <span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">)</span>

import torch

import torch.nn as nn

from typing import Dict, Optional

class KVCache:

def __init__(self, max_batch_size: int, max_seq_len: int,

n_heads: int, head_dim: int, device: str = "cuda"):

self.cache = {

"key": torch.zeros(

(max_batch_size, n_heads, max_seq_len, head_dim),

device=device),

"value": torch.zeros(

device=device)

}

self.position = 0

def update(self,

new_key: torch.Tensor, # (batch, n_heads, new_tokens, head_dim)

new_value: torch.Tensor) -> Dict[str, torch.Tensor]:

batch_size = new_key.size(0)

seq_len = new_key.size(2)

# Update cache slices

self.cache["key"][:batch_size, :, self.position:self.position+seq_len] = new_key

self.cache["value"][:batch_size, :, self.position:self.position+seq_len] = new_value

self.position += seq_len

return {

"key": self.cache["key"][:batch_size, :, :self.position],

"value": self.cache["value"][:batch_size, :, :self.position]

}

class CausalSelfAttention(nn.Module):

def __init__(self, d_model: int, n_heads: int):

super().__init__()

self.d_head = d_model // n_heads

self.n_heads = n_heads

self.W_q = nn.Linear(d_model, d_model)

self.W_k = nn.Linear(d_model, d_model)

self.W_v = nn.Linear(d_model, d_model)

def forward(self,

x: torch.Tensor,

kv_cache: Optional[KVCache] = None) -> torch.Tensor:

batch_size, seq_len, _ = x.shape

# Project to Q/K/V

Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)

K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)

V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)

# Update cache if provided

if kv_cache is not None:

cache = kv_cache.update(K, V)

K, V = cache["key"], cache["value"]

# Scaled dot-product attention

attn = (Q @ K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_head))

mask = torch.tril(torch.ones(seq_len, seq_len)).to(x.device)

attn = attn.masked_fill(mask == 0, float('-inf'))

attn = torch.softmax(attn, dim=-1) @ V

return attn.transpose(1, 2).reshape(batch_size, seq_len, -1)

3. Benchmarking Performance Gains

Test Script

<span class="token keyword">import</span> time
<span class="token keyword">from</span> transformers <span class="token keyword">import</span> AutoModelForCausalLM<span class="token punctuation">,</span> AutoTokenizer

<span class="token keyword">def</span> <span class="token function">benchmark</span><span class="token punctuation">(</span>model<span class="token punctuation">,</span> prompt<span class="token punctuation">,</span> max_new_tokens<span class="token operator">=</span><span class="token number">100</span><span class="token punctuation">,</span> use_cache<span class="token operator">=</span><span class="token boolean">True</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
    inputs <span class="token operator">=</span> tokenizer<span class="token punctuation">(</span>prompt<span class="token punctuation">,</span> return_tensors<span class="token operator">=</span><span class="token string">"pt"</span><span class="token punctuation">)</span><span class="token punctuation">.</span>to<span class="token punctuation">(</span><span class="token string">"cuda"</span><span class="token punctuation">)</span>
    
    start <span class="token operator">=</span> time<span class="token punctuation">.</span>time<span class="token punctuation">(</span><span class="token punctuation">)</span>
    outputs <span class="token operator">=</span> model<span class="token punctuation">.</span>generate<span class="token punctuation">(</span>
        <span class="token operator">**</span>inputs<span class="token punctuation">,</span>
        max_new_tokens<span class="token operator">=</span>max_new_tokens<span class="token punctuation">,</span>
        use_cache<span class="token operator">=</span>use_cache
    <span class="token punctuation">)</span>
    elapsed <span class="token operator">=</span> time<span class="token punctuation">.</span>time<span class="token punctuation">(</span><span class="token punctuation">)</span> <span class="token operator">-</span> start
    
    <span class="token keyword">return</span> elapsed<span class="token punctuation">,</span> tokenizer<span class="token punctuation">.</span>decode<span class="token punctuation">(</span>outputs<span class="token punctuation">[</span><span class="token number">0</span><span class="token punctuation">]</span><span class="token punctuation">)</span>

<span class="token comment"># Load model</span>
model <span class="token operator">=</span> AutoModelForCausalLM<span class="token punctuation">.</span>from_pretrained<span class="token punctuation">(</span><span class="token string">"meta-llama/Llama-2-7b-chat-hf"</span><span class="token punctuation">)</span><span class="token punctuation">.</span>half<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span>cuda<span class="token punctuation">(</span><span class="token punctuation">)</span>

<span class="token comment"># Benchmark</span>
prompt <span class="token operator">=</span> <span class="token string">"Explain quantum computing in simple terms:"</span>
time_cache<span class="token punctuation">,</span> _ <span class="token operator">=</span> benchmark<span class="token punctuation">(</span>model<span class="token punctuation">,</span> prompt<span class="token punctuation">,</span> use_cache<span class="token operator">=</span><span class="token boolean">True</span><span class="token punctuation">)</span>
time_no_cache<span class="token punctuation">,</span> _ <span class="token operator">=</span> benchmark<span class="token punctuation">(</span>model<span class="token punctuation">,</span> prompt<span class="token punctuation">,</span> use_cache<span class="token operator">=</span><span class="token boolean">False</span><span class="token punctuation">)</span>

<span class="token keyword">print</span><span class="token punctuation">(</span><span class="token string-interpolation"><span class="token string">f"With KV Cache: </span><span class="token interpolation"><span class="token punctuation">{</span>time_cache<span class="token punctuation">:</span><span class="token format-spec">.2f</span><span class="token punctuation">}</span></span><span class="token string">s"</span></span><span class="token punctuation">)</span>
<span class="token keyword">print</span><span class="token punctuation">(</span><span class="token string-interpolation"><span class="token string">f"Without KV Cache: </span><span class="token interpolation"><span class="token punctuation">{</span>time_no_cache<span class="token punctuation">:</span><span class="token format-spec">.2f</span><span class="token punctuation">}</span></span><span class="token string">s"</span></span><span class="token punctuation">)</span>
<span class="token keyword">print</span><span class="token punctuation">(</span><span class="token string-interpolation"><span class="token string">f"Speedup: </span><span class="token interpolation"><span class="token punctuation">{</span>time_no_cache<span class="token operator">/</span>time_cache<span class="token punctuation">:</span><span class="token format-spec">.1f</span><span class="token punctuation">}</span></span><span class="token string">x"</span></span><span class="token punctuation">)</span>

import time

from transformers import AutoModelForCausalLM, AutoTokenizer

def benchmark(model, prompt, max_new_tokens=100, use_cache=True):

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

start = time.time()

outputs = model.generate(

**inputs,

max_new_tokens=max_new_tokens,

use_cache=use_cache

)

elapsed = time.time() - start

return elapsed, tokenizer.decode(outputs[0])

# Load model

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf").half().cuda()

# Benchmark

prompt = "Explain quantum computing in simple terms:"

time_cache, _ = benchmark(model, prompt, use_cache=True)

time_no_cache, _ = benchmark(model, prompt, use_cache=False)

print(f"With KV Cache: {time_cache:.2f}s")

print(f"Without KV Cache: {time_no_cache:.2f}s")

print(f"Speedup: {time_no_cache/time_cache:.1f}x")

Results (NVIDIA A100)

Sequence Length	KV Cache (s)	No Cache (s)	Speedup
128 tokens	0.8	3.2	4.0x
512 tokens	2.1	18.7	8.9x
1024 tokens	3.9	67.4	17.3x

4. Advanced Optimizations

Grouped-Query Attention (GQA)

Modern models like Llama-2 use grouped queries to reduce memory overhead:

<span class="token comment"># GQA implementation snippet</span>
num_kv_heads <span class="token operator">=</span> n_heads <span class="token operator">//</span> <span class="token number">8</span>  <span class="token comment"># Group size of 8</span>
K <span class="token operator">=</span> K<span class="token punctuation">.</span>reshape<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> num_kv_heads<span class="token punctuation">,</span> <span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">)</span><span class="token punctuation">.</span>repeat_interleave<span class="token punctuation">(</span><span class="token number">8</span><span class="token punctuation">,</span> dim<span class="token operator">=</span><span class="token number">2</span><span class="token punctuation">)</span>
V <span class="token operator">=</span> V<span class="token punctuation">.</span>reshape<span class="token punctuation">(</span>batch_size<span class="token punctuation">,</span> seq_len<span class="token punctuation">,</span> num_kv_heads<span class="token punctuation">,</span> <span class="token operator">-</span><span class="token number">1</span><span class="token punctuation">)</span><span class="token punctuation">.</span>repeat_interleave<span class="token punctuation">(</span><span class="token number">8</span><span class="token punctuation">,</span> dim<span class="token operator">=</span><span class="token number">2</span><span class="token punctuation">)</span>

# GQA implementation snippet

num_kv_heads = n_heads // 8 # Group size of 8

K = K.reshape(batch_size, seq_len, num_kv_heads, -1).repeat_interleave(8, dim=2)

V = V.reshape(batch_size, seq_len, num_kv_heads, -1).repeat_interleave(8, dim=2)

Memory-Efficient Cache Formats

<span class="token comment"># 4-bit quantized cache</span>
cache <span class="token operator">=</span> torch<span class="token punctuation">.</span>quantize_per_tensor<span class="token punctuation">(</span>
    cache<span class="token punctuation">,</span> scale<span class="token operator">=</span><span class="token number">0.1</span><span class="token punctuation">,</span> zero_point<span class="token operator">=</span><span class="token number">0</span><span class="token punctuation">,</span> dtype<span class="token operator">=</span>torch<span class="token punctuation">.</span>qint4<span class="token punctuation">)</span>

# 4-bit quantized cache

cache = torch.quantize_per_tensor(

cache, scale=0.1, zero_point=0, dtype=torch.qint4)

5. Production Considerations

Best Practices:

Batch Inference: Cache must handle variable-length sequences

<span class="token comment"># Pad sequences to max length in batch</span>
cache <span class="token operator">=</span> torch<span class="token punctuation">.</span>zeros<span class="token punctuation">(</span>batch<span class="token punctuation">,</span> max_len<span class="token punctuation">,</span> <span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">.</span><span class="token punctuation">)</span>

# Pad sequences to max length in batch

cache = torch.zeros(batch, max_len, ...)

Memory Management:

<span class="token comment"># Clear cache between requests</span>
cache<span class="token punctuation">.</span>fill_<span class="token punctuation">(</span><span class="token number">0</span><span class="token punctuation">)</span>  
position <span class="token operator">=</span> <span class="token number">0</span>

# Clear cache between requests

cache.fill_(0)

position = 0

Continuous Batching:

<span class="token comment"># For streaming applications</span>
cache <span class="token operator">=</span> cache<span class="token punctuation">.</span>roll<span class="token punctuation">(</span><span class="token operator">-</span>completed_tokens<span class="token punctuation">,</span> dims<span class="token operator">=</span><span class="token number">1</span><span class="token punctuation">)</span>

# For streaming applications