Building a KV Cache Memory Calculator

Goal

Implement the KV cache memory formula in code, confirm the effect of GQA and quantization by calculation, and work out the maximum number of concurrent requests possible on a given GPU.

Why it matters

The question "Does this model fit on a single A100 80GB?" always comes up before deployment. If you calculate only the weights, the answer is wrong. The weights of a 7B model are 14GB at fp16, but with a 128K context and batch 8, the KV cache is 512GiB. That is 36 times the weights. If you start without this calculation, you meet an OOM on deployment day, and since an OOM occurs when the load piles up, it occurs at the worst possible time. And once you build this calculator, you keep using it — you need it each time you change the model, extend the context length, consider quantization, or decide max_num_seqs. The maximum number of concurrent requests you find in step 6 is the number that becomes the starting point of capacity planning.

Steps

Create kv_bytes(layers, hidden, seq, batch, dtype_bytes, kv_groups=1) in /root/kv/kv.py. The basic formula is 2 * layers * hidden * seq * batch * dtype_bytes.
Write the value calculated with layers=32, hidden=4096, seq=4096, batch=1, dtype_bytes=2 in /root/kv/base.txt as bytes=<정수> gib=<소수 둘째자리> (the placeholders are an integer and a decimal with two decimal places). gib must be 2.00.
Write the value with seq=131072 in /root/kv/long.txt in the same format. gib must be 64.00.
Apply kv_groups=8 and write it in /root/kv/gqa.txt. gib must be 8.00 (on a 128K basis).
Apply dtype_bytes=1 and write it in /root/kv/int8.txt. gib must be 4.00 (on a 128K, GQA-8 basis).
Write gpu_gib=80 weights_gib=14 overhead_gib=6 kv_budget_gib=60 per_req_gib=<소수> max_concurrent=<정수> in /root/kv/fit.txt (the placeholders are a decimal and an integer). The per-request basis is a 32K context, GQA-8 and fp16.
Write the header context,batch,gib and 6 rows in /root/kv/kv_table.csv. The contexts are 4096, 32768 and 131072 and the batches are 1 and 8. The values are on an fp16 basis without GQA.

Notes

The 2 at the front of the formula means two sets, keys and values.
GQA shares key/value heads in groups, so you divide by the number of groups.
Weight quantization (GPTQ, AWQ) and KV cache quantization are separate settings.
Common mistake 1: calculating only the weights and leaving out the KV cache — in long contexts, KV dominates.
Common mistake 2: mixing GiB and GB — this lab is based on 2^30.

Implement the formula

Create kv_bytes(layers, hidden, seq, batch, dtype_bytes, kv_groups=1) in /root/kv/kv.py. The basic formula is 2 * layers * hidden * seq * batch * dtype_bytes.

Put it in after understanding what the 2 at the front is. The arguments are the number of layers, the hidden size, the sequence length, the batch and the dtype bytes.

Verify with the 7B 4K reference value

Write the value calculated with layers=32, hidden=4096, seq=4096, batch=1, dtype_bytes=2 in /root/kv/base.txt as bytes=<정수> gib=<소수 둘째자리> (the placeholders are an integer and a decimal with two decimal places). gib must be 2.00.

There is a widely cited reference value. If it matches here, the formula is right.

Extend to a long context

Write the value with seq=131072 in /root/kv/long.txt in the same format. gib must be 64.00.

Just increase the sequence length by 32 times. See why the result makes long context hard.

Reduce it by applying GQA

Apply kv_groups=8 and write it in /root/kv/gqa.txt. gib must be 8.00 (on a 128K basis).

The key and value heads are shared in groups. Think about where in the formula the division by the number of groups happens.

Apply KV cache quantization

Apply dtype_bytes=1 and write it in /root/kv/int8.txt. gib must be 4.00 (on a 128K, GQA-8 basis).

It is reducing the dtype bytes. What matters is that it is a separate setting from weight quantization.

Find the maximum number of concurrent requests that fit on the GPU

Write gpu_gib=80 weights_gib=14 overhead_gib=6 kv_budget_gib=60 per_req_gib=<소수> max_concurrent=<정수> in /root/kv/fit.txt (the placeholders are a decimal and an integer). The per-request basis is a 32K context, GQA-8 and fp16.

Subtract the weights and overhead from the total memory and divide by the per-request size. This value is the concurrency upper limit.

Generate a scenario table

Write the header context,batch,gib and 6 rows in /root/kv/kv_table.csv. The contexts are 4096, 32768 and 131072 and the batches are 1 and 8. The values are on an fp16 basis without GQA.

Build the table while varying the context and batch. Make it in a form you can use as it is later for capacity planning.