When you need this. A retrieval assistant accepts long documents but runs out of memory only when several readers are active. The model weights fit; the combination of active contexts and runtime workspace does not. Build a byte budget before choosing concurrency limits.
What you run. Executable capacity calculator with explicit example dimensions and workspace assumptions; not a GPU allocation or successful multi-GPU deployment.
This lesson builds on benchmark streaming inference under load.
Count bytes before choosing concurrency
For a dense attention cache, each stored token needs both keys and values for every layer. Its payload is 2 × layers × KV heads × head dimension × bytes per element. Grouped-query attention (GQA) shares keys and values across query heads; use KV heads, not query heads, in that formula.
Add weight bytes, cache bytes, and explicit workspace reserves in one unit. GB means decimal bytes; GiB means binary bytes. Likewise, KiB and MiB use binary units for per-token payloads. The calculator now keeps these units separate. An aggregate byte budget is only a lower bound: tensor-parallel sharding, replicated KV heads, alignment, and per-device workspace can still prevent a configuration from fitting.
size_inference_memory() computes the cache payload and reports decimal totals alongside clearly labeled binary quantities. Its example model dimensions are calculator fixtures. Compare them with the actual small-model configuration recorded in the first lesson; they are different workloads.
Run the memory budget
Download the lab from this page, extract it, and run these commands from its directory. The default command reads saved run records or runs the explicitly described local algorithm; it never downloads a model or calls a paid API.
python3 lab.py
python3 -m unittest -v test_lab.pyRead lab.py
"""GPU memory calculator for model weights, MHA vs GQA KV cache, activations, and fleet sizing."""
from __future__ import annotations
from dataclasses import asdict, dataclass
import json
import math
@dataclass(frozen=True)
class ModelConfig:
name: str
params_billions: float
num_layers: int
num_query_heads: int
num_kv_heads: int
head_dim: int
weight_bytes: int = 2
kv_bytes: int = 2
@dataclass(frozen=True)
class MemoryBreakdown:
config_name: str
seq_len: int
batch_size: int
weights_gb: float
kv_bytes_per_token: int
kv_kib_per_token: float
kv_mib_per_token: float
kv_cache_gib: float
kv_cache_gb: float
activation_and_cuda_gb: float
total_required_gb: float
min_h100_80gb_gpus: int
min_h200_141gb_gpus: int
fits_single_8x_h100_node: bool
def size_inference_memory(
config: ModelConfig,
seq_len: int = 32768,
batch_size: int = 100,
include_activation_headroom: bool = True,
) -> MemoryBreakdown:
if seq_len <= 0 or batch_size <= 0:
raise ValueError("seq_len and batch_size must be positive")
# 1. Model weights: 70B parameters * bytes_per_param (e.g., 140 GB in BF16, 70 GB in FP8).
weights_gb = round(config.params_billions * config.weight_bytes, 2)
# 2. Exact KV cache per token: 2 (K and V) * layers * kv_heads * head_dim * kv_bytes.
kv_bytes_per_token = (
2 * config.num_layers * config.num_kv_heads * config.head_dim * config.kv_bytes
)
kv_kib_per_token = round(kv_bytes_per_token / 1024.0, 2)
kv_mib_per_token = round(kv_bytes_per_token / (1024.0 * 1024.0), 4)
total_kv_bytes = kv_bytes_per_token * seq_len * batch_size
kv_cache_gib = round(total_kv_bytes / (1024.0**3), 2)
kv_cache_gb = round(total_kv_bytes / 1e9, 2)
# 3. Prefill activation workspace + CUDA context / NCCL communication buffers.
if include_activation_headroom:
prefill_workspace_gb = round(4.5 + 0.00035 * seq_len, 2)
cuda_and_nccl_gb = 3.5
activation_and_cuda_gb = round(prefill_workspace_gb + cuda_and_nccl_gb, 2)
else:
activation_and_cuda_gb = 0.0
total_required_gb = round(
weights_gb + total_kv_bytes / 1e9 + activation_and_cuda_gb, 2
)
min_h100 = max(1, int(math.ceil(total_required_gb / 80.0)))
min_h200 = max(1, int(math.ceil(total_required_gb / 141.0)))
return MemoryBreakdown(
config_name=config.name,
seq_len=seq_len,
batch_size=batch_size,
weights_gb=weights_gb,
kv_bytes_per_token=kv_bytes_per_token,
kv_kib_per_token=kv_kib_per_token,
kv_mib_per_token=kv_mib_per_token,
kv_cache_gib=kv_cache_gib,
kv_cache_gb=kv_cache_gb,
activation_and_cuda_gb=activation_and_cuda_gb,
total_required_gb=total_required_gb,
min_h100_80gb_gpus=min_h100,
min_h200_141gb_gpus=min_h200,
fits_single_8x_h100_node=(min_h100 <= 8),
)
def validate_node_deployment(
breakdown: MemoryBreakdown,
actual_activation_gb: float = 19.47,
node_vram_gb: float = 640.0,
) -> None:
if (
not breakdown.fits_single_8x_h100_node
or breakdown.total_required_gb > node_vram_gb
):
raise MemoryError(
f"{breakdown.config_name} requires {breakdown.total_required_gb} GB, exceeding 8xH100 node ({node_vram_gb} GB)"
)
if breakdown.activation_and_cuda_gb < actual_activation_gb:
peak_gb = round(
breakdown.total_required_gb
- breakdown.activation_and_cuda_gb
+ actual_activation_gb,
2,
)
if peak_gb > node_vram_gb:
raise MemoryError(
f"Prefill OOM: unbudgeted activations ({actual_activation_gb} GB) push peak to {peak_gb} GB > {node_vram_gb} GB"
)
def main() -> None:
mha_bf16 = ModelConfig(
name="llama70b-mha-bf16",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=64,
head_dim=128,
weight_bytes=2,
kv_bytes=2,
)
gqa_bf16 = ModelConfig(
name="llama70b-gqa-bf16",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=8,
head_dim=128,
weight_bytes=2,
kv_bytes=2,
)
gqa_fp8 = ModelConfig(
name="llama70b-gqa-fp8",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=8,
head_dim=128,
weight_bytes=1,
kv_bytes=1,
)
mha_res = size_inference_memory(mha_bf16, seq_len=32768, batch_size=100)
gqa_bf16_res = size_inference_memory(gqa_bf16, seq_len=32768, batch_size=100)
gqa_fp8_res = size_inference_memory(gqa_fp8, seq_len=32768, batch_size=100)
report = {
"experiment": "gpu-memory-sizing-weights-and-kv-cache",
"workload": {"seq_len": 32768, "concurrent_users": 100},
"configurations": {
"mha_bf16_64kv": asdict(mha_res),
"gqa_bf16_8kv": asdict(gqa_bf16_res),
"gqa_fp8_8kv": asdict(gqa_fp8_res),
},
"reduction_factors": {
"gqa_vs_mha_kv_ratio": round(
mha_res.kv_bytes_per_token / gqa_bf16_res.kv_bytes_per_token, 1
),
"gqa_fp8_vs_mha_bf16_kv_ratio": round(
mha_res.kv_bytes_per_token / gqa_fp8_res.kv_bytes_per_token, 1
),
},
}
print(
json.dumps(
{
"execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
**report,
},
indent=2,
)
)
if __name__ == "__main__":
main()The binary and decimal cache figures describe the same payload. Adding a GiB value directly to a GB value would understate the total. The aggregate device count is a capacity lower bound; running that sharding layout still requires validation.
| Quantity | Recorded value |
|---|---|
| Weight payload (GB) | 140 |
| GQA cache payload (GiB) | 1000 |
| GQA cache payload (GB) | 1073.74 |
| Total example budget (GB) | 1233.21 |
Inspect the complete local result
{
"execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
"experiment": "gpu-memory-sizing-weights-and-kv-cache",
"workload": {
"seq_len": 32768,
"concurrent_users": 100
},
"configurations": {
"mha_bf16_64kv": {
"config_name": "llama70b-mha-bf16",
"seq_len": 32768,
"batch_size": 100,
"weights_gb": 140.0,
"kv_bytes_per_token": 2621440,
"kv_kib_per_token": 2560.0,
"kv_mib_per_token": 2.5,
"kv_cache_gib": 8000.0,
"kv_cache_gb": 8589.93,
"activation_and_cuda_gb": 19.47,
"total_required_gb": 8749.4,
"min_h100_80gb_gpus": 110,
"min_h200_141gb_gpus": 63,
"fits_single_8x_h100_node": false
},
"gqa_bf16_8kv": {
"config_name": "llama70b-gqa-bf16",
"seq_len": 32768,
"batch_size": 100,
"weights_gb": 140.0,
"kv_bytes_per_token": 327680,
"kv_kib_per_token": 320.0,
"kv_mib_per_token": 0.3125,
"kv_cache_gib": 1000.0,
"kv_cache_gb": 1073.74,
"activation_and_cuda_gb": 19.47,
"total_required_gb": 1233.21,
"min_h100_80gb_gpus": 16,
"min_h200_141gb_gpus": 9,
"fits_single_8x_h100_node": false
},
"gqa_fp8_8kv": {
"config_name": "llama70b-gqa-fp8",
"seq_len": 32768,
"batch_size": 100,
"weights_gb": 70.0,
"kv_bytes_per_token": 163840,
"kv_kib_per_token": 160.0,
"kv_mib_per_token": 0.1562,
"kv_cache_gib": 500.0,
"kv_cache_gb": 536.87,
"activation_and_cuda_gb": 19.47,
"total_required_gb": 626.34,
"min_h100_80gb_gpus": 8,
"min_h200_141gb_gpus": 5,
"fits_single_8x_h100_node": true
}
},
"reduction_factors": {
"gqa_vs_mha_kv_ratio": 8.0,
"gqa_fp8_vs_mha_bf16_kv_ratio": 16.0
}
}
This image is an application-generated report of the saved program result, not a new inference run. Open the full-resolution report or read the raw result.
Leave room for peak execution
Lower cache precision can increase capacity, but requires supported kernels, scale storage, and quality checks. Reserving more space for active contexts leaves less room for prefill workspace and graph buffers. Planning from average context length misses a burst of unusually long documents.
Insufficient aggregate capacity. The example byte budget exceeds the supplied node capacity, so the calculator refuses that plan. It does not attempt a GPU allocation.
Missing workspace reserve. A plan that fits without its workspace reserve can exceed the same capacity when that reserve is included; the validator rejects it.
Read the executable checks
"""Unit and failure-boundary tests for the GPU memory and KV cache calculator."""
from __future__ import annotations
import unittest
from lab import ModelConfig, size_inference_memory, validate_node_deployment
class GpuMemorySizingTests(unittest.TestCase):
def test_exact_kv_cache_formulas_for_mha_and_gqa(self) -> None:
mha = ModelConfig(
name="mha-bf16",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=64,
head_dim=128,
weight_bytes=2,
kv_bytes=2,
)
gqa_bf16 = ModelConfig(
name="gqa-bf16",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=8,
head_dim=128,
weight_bytes=2,
kv_bytes=2,
)
gqa_fp8 = ModelConfig(
name="gqa-fp8",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=8,
head_dim=128,
weight_bytes=1,
kv_bytes=1,
)
mha_res = size_inference_memory(mha, seq_len=32768, batch_size=100)
gqa_bf16_res = size_inference_memory(gqa_bf16, seq_len=32768, batch_size=100)
gqa_fp8_res = size_inference_memory(gqa_fp8, seq_len=32768, batch_size=100)
self.assertEqual(mha_res.kv_mib_per_token, 2.5)
self.assertEqual(mha_res.kv_cache_gib, 8000.0)
self.assertEqual(gqa_bf16_res.kv_kib_per_token, 320.0)
self.assertEqual(gqa_bf16_res.kv_cache_gb, 1073.74)
self.assertEqual(gqa_fp8_res.kv_kib_per_token, 160.0)
self.assertEqual(gqa_fp8_res.kv_cache_gb, 536.87)
self.assertTrue(gqa_fp8_res.fits_single_8x_h100_node)
def test_mha_kv_cache_overflows_8x_h100_node(self) -> None:
mha = ModelConfig(
name="llama70b-mha-bf16",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=64,
head_dim=128,
weight_bytes=2,
kv_bytes=2,
)
breakdown = size_inference_memory(mha, seq_len=32768, batch_size=100)
self.assertFalse(breakdown.fits_single_8x_h100_node)
self.assertGreater(breakdown.min_h100_80gb_gpus, 100)
with self.assertRaisesRegex(MemoryError, "exceeding 8xH100 node"):
validate_node_deployment(breakdown, node_vram_gb=640.0)
def test_ignoring_activation_headroom_causes_prefill_oom(self) -> None:
gqa_fp8 = ModelConfig(
name="llama70b-gqa-fp8-tight",
params_billions=70.0,
num_layers=80,
num_query_heads=64,
num_kv_heads=8,
head_dim=128,
weight_bytes=1,
kv_bytes=1,
)
# Decimal GB throughout: the unbudgeted total fits, but the explicit workspace reserve does not.
unbudgeted = size_inference_memory(
gqa_fp8,
seq_len=32768,
batch_size=104,
include_activation_headroom=False,
)
self.assertLessEqual(unbudgeted.total_required_gb, 640.0)
with self.assertRaisesRegex(MemoryError, "Prefill OOM: unbudgeted activations"):
validate_node_deployment(
unbudgeted, actual_activation_gb=19.47, node_vram_gb=640.0
)
class DecimalBudgetTests(unittest.TestCase):
def test_total_budget_adds_decimal_gb(self):
model = ModelConfig("example", 70, 80, 64, 8, 128)
result = size_inference_memory(model)
self.assertEqual(result.total_required_gb, 1233.21)
self.assertEqual(result.kv_cache_gb, 1073.74)
if __name__ == "__main__":
unittest.main()
class GuardBoundaryTests(unittest.TestCase):
def test_rejects_invalid_boundary_inputs_with_specific_reason(self):
import lab
cases = [('size_inference_memory(ModelConfig("test",1,1,1,1,1),seq_len=0)', 'seq_len and')]
for code, reason in cases:
with self.subTest(code=code):
with self.assertRaisesRegex((ValueError, MemoryError, KeyError, RuntimeError), reason):
exec(code, vars(lab).copy())Try a change and explain the result
Change the active sequence count while holding context length fixed. Then change the number of KV heads. Predict which terms change and verify the byte arithmetic. Explain why adding total memory across several devices does not prove a valid sharding layout.
Read the explanation
Stored key/value tensors use KV heads. Query-head count can overestimate cache payload in a GQA configuration.
Validate the target memory layout
Inspect the actual model config and engine startup allocation. Validate a long-context, high-concurrency workload and leave measured workspace headroom. Use peak allocated and reserved memory from the target runtime rather than a universal percentage.
References
Sources checked October 24, 2026. Runtime-specific commands and results apply to the recorded versions; linked rolling documentation can change. The application scenarios are teaching examples, not claims about an author’s production incident.