Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Mistral Architecture › Rolling Buffer KV Cache

Michael BrenndoerferUpdated August 6, 20255 min read

Language AI HandbookAll chapters

Part 29 of 67Modern Decoder Models

Chapter 211Mistral Architecture › Rolling Buffer KV Cache

Reading section 4 of 14 · Rolling Buffer KV Cache · 5 minSection 4 of 14 · 5 min

Section 4 of 14 · 5 min read

Mistral Architecture › Rolling Buffer KV CacheOne section · Sections in this chapter4 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

During autoregressive generation, transformers cache the key and value projections of previous tokens to avoid redundant computation. In standard transformers, this KV cache grows linearly with sequence length, eventually consuming significant memory for long generations.

Mistral implements a rolling buffer KV cache that exploits sliding window attention. Since tokens beyond the window boundary will never be attended to again, their cached keys and values can be safely discarded. The cache maintains only the most recent $w$ positions, using modular indexing to overwrite old entries.

In[13]:

Code

class RollingKVCache:
    """
    Rolling buffer cache for sliding window attention.

    Only stores the most recent `window_size` key-value pairs,
    using modular indexing to overwrite old entries.
    """

    def __init__(self, window_size, num_heads, head_dim, dtype=np.float16):
        self.window_size = window_size
        self.num_heads = num_heads
        self.head_dim = head_dim

        # Pre-allocate fixed-size buffers
        self.key_cache = np.zeros(
            (window_size, num_heads, head_dim), dtype=dtype
        )
        self.value_cache = np.zeros(
            (window_size, num_heads, head_dim), dtype=dtype
        )
        self.current_length = 0

    def update(self, position, key, value):
        """Add new key-value pair at the given position."""
        # Use modular indexing to wrap around
        cache_idx = position % self.window_size
        self.key_cache[cache_idx] = key
        self.value_cache[cache_idx] = value
        self.current_length = min(self.current_length + 1, self.window_size)

    def get_valid_entries(self, current_position):
        """Retrieve cached entries visible from current position."""
        if current_position < self.window_size:
            # Haven't filled the buffer yet
            return self.key_cache[: current_position + 1], self.value_cache[
                : current_position + 1
            ]
        else:
            # Buffer is full, need to handle wrap-around
            start_pos = current_position - self.window_size + 1
            indices = [
                (start_pos + i) % self.window_size
                for i in range(self.window_size)
            ]
            return self.key_cache[indices], self.value_cache[indices]

    def memory_usage_bytes(self):
        """Calculate actual memory usage."""
        return self.key_cache.nbytes + self.value_cache.nbytes

Out[14]:

Console

Rolling KV Cache Demonstration (window=4):

As we add tokens, old entries are overwritten:

Position 0: Stored at index 0, visible range [0, 0]
Position 1: Stored at index 1, visible range [0, 1]
Position 2: Stored at index 2, visible range [0, 2]
Position 3: Stored at index 3, visible range [0, 3]
Position 4: Stored at index 0, visible range [1, 4]
Position 5: Stored at index 1, visible range [2, 5]
Position 6: Stored at index 2, visible range [3, 6]
Position 7: Stored at index 3, visible range [4, 7]

The rolling buffer ensures that memory usage remains constant regardless of how many tokens are generated. For a 32-layer Mistral model with 8 KV heads per layer, the total cache memory is fixed at:

\text{Cache Size} = 2 \times L \times w \times h_{\text{kv}} \times d_h \times b

where:

$\text{Cache Size}$ : total memory in bytes for the KV cache
$2$ : factor accounting for storing both keys and values (each cached position stores one K vector and one V vector)
$L$ : number of transformer layers (32 for Mistral 7B), since each layer maintains its own KV cache
$w$ : sliding window size (4096 for Mistral), the maximum number of positions stored
$h_{\text{kv}}$ : number of key-value heads per layer (8 for Mistral 7B with GQA)
$d_h$ : dimension of each head, computed as $d_h = d_{\text{model}} / h_q$ where $d_{\text{model}} = 4096$ and $h_q = 32$ , giving $d_h = 128$
$b$ : bytes per element (2 for fp16, 4 for fp32)

For Mistral 7B with fp16 precision, the fixed cache size is:

2 \times 32 \times 4096 \times 8 \times 128 \times 2 = 536,870,912 \text{ bytes} \approx 512 \text{ MB}

This fixed 512 MB footprint contrasts sharply with standard transformers, where the cache grows linearly with sequence length and could reach several gigabytes for long generations.

In[15]:

Code

def kv_cache_size_bytes(
    num_layers, window_size, num_kv_heads, head_dim, dtype_bytes=2
):
    """Calculate KV cache memory for rolling buffer."""
    # 2 for keys and values
    return 2 * num_layers * window_size * num_kv_heads * head_dim * dtype_bytes


def compare_cache_sizes(seq_lengths, config):
    """Compare rolling vs standard KV cache."""
    num_layers = config["num_layers"]
    num_kv_heads = config["num_kv_heads"]
    head_dim = config["hidden_size"] // config["num_attention_heads"]
    window_size = config["sliding_window"]

    results = []
    for seq_len in seq_lengths:
        # Standard cache: stores all positions
        standard_size = kv_cache_size_bytes(
            num_layers, seq_len, num_kv_heads, head_dim
        )
        # Rolling cache: stores only window_size positions
        rolling_size = kv_cache_size_bytes(
            num_layers, min(seq_len, window_size), num_kv_heads, head_dim
        )
        results.append(
            {
                "seq_len": seq_len,
                "standard_mb": standard_size / (1024**2),
                "rolling_mb": rolling_size / (1024**2),
            }
        )
    return results


cache_comparison = compare_cache_sizes(seq_lengths, mistral_config)

Out[16]:

Console

KV Cache Memory Comparison (Mistral 7B config):

 Sequence Length   Standard Cache    Rolling Cache
----------------------------------------------------
         4,096.0          512.0 MB          512.0 MB
4,404.743718592965          550.6 MB          512.0 MB
4,713.48743718593          589.2 MB          512.0 MB
5,022.231155778894          627.8 MB          512.0 MB
5,330.974874371859          666.4 MB          512.0 MB
5,639.718592964824          705.0 MB          512.0 MB
5,948.462311557789          743.6 MB          512.0 MB
6,257.206030150754          782.2 MB          512.0 MB
6,565.949748743718          820.7 MB          512.0 MB
6,874.693467336683          859.3 MB          512.0 MB
7,183.437185929648          897.9 MB          512.0 MB
7,492.180904522613          936.5 MB          512.0 MB
7,800.924623115578          975.1 MB          512.0 MB
8,109.668341708542         1013.7 MB          512.0 MB
8,418.412060301507         1052.3 MB          512.0 MB
8,727.155778894472         1090.9 MB          512.0 MB
9,035.899497487437         1129.5 MB          512.0 MB
9,344.643216080402         1168.1 MB          512.0 MB
9,653.386934673366         1206.7 MB          512.0 MB
9,962.130653266331         1245.3 MB          512.0 MB
10,270.874371859296         1283.9 MB          512.0 MB
10,579.61809045226         1322.5 MB          512.0 MB
10,888.361809045226         1361.0 MB          512.0 MB
11,197.10552763819         1399.6 MB          512.0 MB
11,505.849246231155         1438.2 MB          512.0 MB
11,814.59296482412         1476.8 MB          512.0 MB
12,123.336683417085         1515.4 MB          512.0 MB
12,432.08040201005         1554.0 MB          512.0 MB
12,740.824120603014         1592.6 MB          512.0 MB
13,049.56783919598         1631.2 MB          512.0 MB
13,358.311557788944         1669.8 MB          512.0 MB
13,667.055276381909         1708.4 MB          512.0 MB
13,975.798994974873         1747.0 MB          512.0 MB
14,284.542713567838         1785.6 MB          512.0 MB
14,593.286432160803         1824.2 MB          512.0 MB
14,902.030150753768         1862.8 MB          512.0 MB
15,210.773869346733         1901.3 MB          512.0 MB
15,519.517587939697         1939.9 MB          512.0 MB
15,828.261306532662         1978.5 MB          512.0 MB
16,137.005025125627         2017.1 MB          512.0 MB
16,445.748743718592         2055.7 MB          512.0 MB
16,754.49246231156         2094.3 MB          512.0 MB
17,063.23618090452         2132.9 MB          512.0 MB
17,371.979899497484         2171.5 MB          512.0 MB
17,680.72361809045         2210.1 MB          512.0 MB
17,989.467336683418         2248.7 MB          512.0 MB
18,298.21105527638         2287.3 MB          512.0 MB
18,606.954773869344         2325.9 MB          512.0 MB
18,915.69849246231         2364.5 MB          512.0 MB
19,224.442211055277         2403.1 MB          512.0 MB
19,533.18592964824         2441.6 MB          512.0 MB
19,841.929648241203         2480.2 MB          512.0 MB
20,150.67336683417         2518.8 MB          512.0 MB
20,459.417085427136         2557.4 MB          512.0 MB
20,768.1608040201         2596.0 MB          512.0 MB
21,076.904522613062         2634.6 MB          512.0 MB
21,385.64824120603         2673.2 MB          512.0 MB
21,694.391959798995         2711.8 MB          512.0 MB
22,003.13567839196         2750.4 MB          512.0 MB
22,311.87939698492         2789.0 MB          512.0 MB
22,620.623115577888         2827.6 MB          512.0 MB
22,929.366834170854         2866.2 MB          512.0 MB
23,238.110552763817         2904.8 MB          512.0 MB
23,546.85427135678         2943.4 MB          512.0 MB
23,855.597989949747         2981.9 MB          512.0 MB
24,164.341708542714         3020.5 MB          512.0 MB
24,473.085427135677         3059.1 MB          512.0 MB
24,781.82914572864         3097.7 MB          512.0 MB
25,090.572864321606         3136.3 MB          512.0 MB
25,399.316582914573         3174.9 MB          512.0 MB
25,708.060301507536         3213.5 MB          512.0 MB
26,016.8040201005         3252.1 MB          512.0 MB
26,325.547738693465         3290.7 MB          512.0 MB
26,634.291457286432         3329.3 MB          512.0 MB
26,943.035175879395         3367.9 MB          512.0 MB
27,251.778894472358         3406.5 MB          512.0 MB
27,560.522613065325         3445.1 MB          512.0 MB
27,869.26633165829         3483.7 MB          512.0 MB
28,178.010050251254         3522.3 MB          512.0 MB
28,486.753768844217         3560.8 MB          512.0 MB
28,795.497487437184         3599.4 MB          512.0 MB
29,104.24120603015         3638.0 MB          512.0 MB
29,412.984924623113         3676.6 MB          512.0 MB
29,721.728643216076         3715.2 MB          512.0 MB
30,030.472361809043         3753.8 MB          512.0 MB
30,339.21608040201         3792.4 MB          512.0 MB
30,647.959798994973         3831.0 MB          512.0 MB
30,956.703517587935         3869.6 MB          512.0 MB
31,265.447236180902         3908.2 MB          512.0 MB
31,574.19095477387         3946.8 MB          512.0 MB
31,882.93467336683         3985.4 MB          512.0 MB
32,191.678391959795         4024.0 MB          512.0 MB
32,500.42211055276         4062.6 MB          512.0 MB
32,809.16582914573         4101.1 MB          512.0 MB
33,117.90954773869         4139.7 MB          512.0 MB
33,426.653266331654         4178.3 MB          512.0 MB
33,735.39698492462         4216.9 MB          512.0 MB
34,044.14070351759         4255.5 MB          512.0 MB
34,352.884422110554         4294.1 MB          512.0 MB
34,661.62814070351         4332.7 MB          512.0 MB
34,970.37185929648         4371.3 MB          512.0 MB
35,279.115577889446         4409.9 MB          512.0 MB
35,587.859296482406         4448.5 MB          512.0 MB
35,896.60301507537         4487.1 MB          512.0 MB
36,205.34673366834         4525.7 MB          512.0 MB
36,514.090452261305         4564.3 MB          512.0 MB
36,822.83417085427         4602.9 MB          512.0 MB
37,131.57788944723         4641.4 MB          512.0 MB
37,440.3216080402         4680.0 MB          512.0 MB
37,749.065326633165         4718.6 MB          512.0 MB
38,057.809045226124         4757.2 MB          512.0 MB
38,366.55276381909         4795.8 MB          512.0 MB
38,675.29648241206         4834.4 MB          512.0 MB
38,984.040201005024         4873.0 MB          512.0 MB
39,292.78391959799         4911.6 MB          512.0 MB
39,601.52763819095         4950.2 MB          512.0 MB
39,910.27135678392         4988.8 MB          512.0 MB
40,219.01507537688         5027.4 MB          512.0 MB
40,527.75879396984         5066.0 MB          512.0 MB
40,836.50251256281         5104.6 MB          512.0 MB
41,145.246231155776         5143.2 MB          512.0 MB
41,453.98994974874         5181.7 MB          512.0 MB
41,762.73366834171         5220.3 MB          512.0 MB
42,071.47738693467         5258.9 MB          512.0 MB
42,380.221105527635         5297.5 MB          512.0 MB
42,688.9648241206         5336.1 MB          512.0 MB
42,997.70854271356         5374.7 MB          512.0 MB
43,306.45226130653         5413.3 MB          512.0 MB
43,615.195979899494         5451.9 MB          512.0 MB
43,923.93969849246         5490.5 MB          512.0 MB
44,232.68341708543         5529.1 MB          512.0 MB
44,541.42713567839         5567.7 MB          512.0 MB
44,850.17085427135         5606.3 MB          512.0 MB
45,158.91457286432         5644.9 MB          512.0 MB
45,467.65829145728         5683.5 MB          512.0 MB
45,776.402010050246         5722.1 MB          512.0 MB
46,085.14572864321         5760.6 MB          512.0 MB
46,393.88944723618         5799.2 MB          512.0 MB
46,702.633165829146         5837.8 MB          512.0 MB
47,011.376884422105         5876.4 MB          512.0 MB
47,320.12060301507         5915.0 MB          512.0 MB
47,628.86432160804         5953.6 MB          512.0 MB
47,937.608040201         5992.2 MB          512.0 MB
48,246.351758793964         6030.8 MB          512.0 MB
48,555.09547738693         6069.4 MB          512.0 MB
48,863.8391959799         6108.0 MB          512.0 MB
49,172.582914572864         6146.6 MB          512.0 MB
49,481.32663316582         6185.2 MB          512.0 MB
49,790.07035175879         6223.8 MB          512.0 MB
50,098.81407035176         6262.4 MB          512.0 MB
50,407.557788944716         6300.9 MB          512.0 MB
50,716.30150753768         6339.5 MB          512.0 MB
51,025.04522613065         6378.1 MB          512.0 MB
51,333.788944723616         6416.7 MB          512.0 MB
51,642.53266331658         6455.3 MB          512.0 MB
51,951.27638190954         6493.9 MB          512.0 MB
52,260.02010050251         6532.5 MB          512.0 MB
52,568.763819095475         6571.1 MB          512.0 MB
52,877.507537688434         6609.7 MB          512.0 MB
53,186.2512562814         6648.3 MB          512.0 MB
53,494.99497487437         6686.9 MB          512.0 MB
53,803.738693467334         6725.5 MB          512.0 MB
54,112.4824120603         6764.1 MB          512.0 MB
54,421.22613065326         6802.7 MB          512.0 MB
54,729.96984924623         6841.2 MB          512.0 MB
55,038.71356783919         6879.8 MB          512.0 MB
55,347.45728643215         6918.4 MB          512.0 MB
55,656.20100502512         6957.0 MB          512.0 MB
55,964.944723618086         6995.6 MB          512.0 MB
56,273.68844221105         7034.2 MB          512.0 MB
56,582.43216080402         7072.8 MB          512.0 MB
56,891.17587939698         7111.4 MB          512.0 MB
57,199.919597989945         7150.0 MB          512.0 MB
57,508.66331658291         7188.6 MB          512.0 MB
57,817.40703517587         7227.2 MB          512.0 MB
58,126.15075376884         7265.8 MB          512.0 MB
58,434.894472361804         7304.4 MB          512.0 MB
58,743.63819095477         7343.0 MB          512.0 MB
59,052.38190954774         7381.5 MB          512.0 MB
59,361.1256281407         7420.1 MB          512.0 MB
59,669.86934673366         7458.7 MB          512.0 MB
59,978.61306532663         7497.3 MB          512.0 MB
60,287.35678391959         7535.9 MB          512.0 MB
60,596.100502512556         7574.5 MB          512.0 MB
60,904.84422110552         7613.1 MB          512.0 MB
61,213.58793969849         7651.7 MB          512.0 MB
61,522.331658291456         7690.3 MB          512.0 MB
61,831.075376884415         7728.9 MB          512.0 MB
62,139.81909547738         7767.5 MB          512.0 MB
62,448.56281407035         7806.1 MB          512.0 MB
62,757.30653266331         7844.7 MB          512.0 MB
63,066.050251256274         7883.3 MB          512.0 MB
63,374.79396984924         7921.8 MB          512.0 MB
63,683.53768844221         7960.4 MB          512.0 MB
63,992.281407035174         7999.0 MB          512.0 MB
64,301.02512562813         8037.6 MB          512.0 MB
64,609.7688442211         8076.2 MB          512.0 MB
64,918.51256281407         8114.8 MB          512.0 MB
65,227.256281407026         8153.4 MB          512.0 MB
        65,536.0         8192.0 MB          512.0 MB

Rolling cache is fixed at 512.0 MB regardless of sequence length

The fixed memory footprint is valuable for deployment scenarios where memory budgets are tight and generation lengths are unpredictable.

Out[17]:

Visualization

Line plot showing standard KV cache growing linearly while rolling buffer stays flat at 512 MB. — KV cache memory comparison between standard (growing) and rolling buffer (fixed) approaches. The rolling buffer caps memory at 512 MB regardless of sequence length, while standard cache grows linearly.