Memory & FLOPs Accounting · AI infra interview lesson

Why It Matters

A 2×3 BF16 matrix has six elements at two bytes each: 12 bytes for weights. Training also stores gradients and optimizer states. Add each item before estimating device capacity.

One concrete number: LLaMA-2-7B (32 layers, 32 KV heads, head_dim=128, FP16) stores 2 * 32 * 32 * 128 * 2 = 524,288 bytes of KV cache per token, exactly 512 KiB. A 4096-token context is 2 GiB, and 50 such concurrent requests need 100 GiB, which does not fit on an 80 GB card. Llama-3-8B uses 8 KV heads and drops this to 128 KiB per token. KV cache management in inference engines, GQA, and KV quantization all start from this calculation.

Walk the Shapes First

A tiny Llama: V=16, d=8, L=2, H=2, H_kv=1, head_dim=4, ffn=12. nn.Linear weights have shape (out, in).

Now one decode step (B=1, 3 tokens of context, generating the 4th):

Rule read off the table: parameter count only depends on weight shapes; K/V projections and the KV cache scale with H_kv, not H; a decode step takes a single token as input but must read every weight and the whole KV history.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: parameter count with GQA and tying

What's new: write the parameter count as a function of the config and check it against numel of a real nn.Module.

Per layer: attn = d*H*hd + 2*d*H_kv*hd + H*hd*d, mlp = 3*d*ffn, norm = 2*d; total V*d*(1 if tie else 2) + L*(attn+mlp+norm) + d. Pass hd explicitly: Qwen3 uses head_dim=128, which need not equal d/H.

import torch
import torch.nn as nn
torch.manual_seed(0)

def count_params(V, d, L, H, H_kv, hd, ffn, tie=False):
    attn = d * H * hd + 2 * d * H_kv * hd + H * hd * d
    mlp = 3 * d * ffn
    norms = 2 * d
    emb = V * d * (1 if tie else 2)
    return emb + L * (attn + mlp + norms) + d

class Layer(nn.Module):
    def __init__(self, d, H, H_kv, hd, ffn):
        super().__init__()
        self.input_layernorm = nn.Parameter(torch.ones(d))
        self.q_proj = nn.Linear(d, H * hd, bias=False)
        self.k_proj = nn.Linear(d, H_kv * hd, bias=False)
        self.v_proj = nn.Linear(d, H_kv * hd, bias=False)
        self.o_proj = nn.Linear(H * hd, d, bias=False)
        self.post_attention_layernorm = nn.Parameter(torch.ones(d))
        self.gate_proj = nn.Linear(d, ffn, bias=False)
        self.up_proj = nn.Linear(d, ffn, bias=False)
        self.down_proj = nn.Linear(ffn, d, bias=False)

Check parameter counts with a tiny model

Count with a formula first, then compare with parameter tensors. Shared weights count once; architecture and dtype are inputs to the calculation.

class TinyLlama(nn.Module):
    def __init__(self, V, d, L, H, H_kv, hd, ffn, tie=False):
        super().__init__()
        self.embed_tokens = nn.Embedding(V, d)
        self.layers = nn.ModuleList(Layer(d, H, H_kv, hd, ffn) for _ in range(L))
        self.norm = nn.Parameter(torch.ones(d))
        self.lm_head = nn.Linear(d, V, bias=False)
        if tie:
            self.lm_head.weight = self.embed_tokens.weight

for tie in (False, True):
    m = TinyLlama(16, 8, 2, 2, 1, 4, 12, tie)
    n_module = sum(p.numel() for p in m.parameters())  # parameters() deduplicates shared weights
    assert n_module == count_params(16, 8, 2, 2, 1, 4, 12, tie)
    print(tie, n_module)        # False 1256 / True 1128

# head_dim != d/H (Qwen3 style)
m = TinyLlama(16, 8, 1, 2, 1, 6, 12)
assert sum(p.numel() for p in m.parameters()) == count_params(16, 8, 1, 2, 1, 6, 12)

# LLaMA-2-7B and Llama-3-8B
assert count_params(32000, 4096, 32, 32, 32, 128, 11008) == 6_738_415_616
assert count_params(128256, 4096, 32, 32, 8, 128, 14336) == 8_030_261_248

model.parameters() deduplicates shared tensors, so after tying the module count also drops by exactly V*d. If you count via state_dict(), two keys point to the same storage and you count it twice.

Step 2: static training memory, 16Ψ

What's new: convert parameter count Ψ into bytes for the standard AdamW + mixed-precision setup.

BF16 weight 2 + BF16 grad 2 + FP32 master weight 4 + Adam m 4 + Adam v 4 = 16 bytes/param. This is only the static part: no activations, temporary buffers, communication buffers, or fragmentation.

import torch
torch.manual_seed(0)

def static_train_bytes(n_params, w=2, g=2, master=4, n_states=2, state_bytes=4):
    return n_params * (w + g + master + n_states * state_bytes)

# allocate real tensors and count bytes
shapes = [(16, 8), (8, 8), (4, 8), (4, 8), (8, 8), (12, 8), (12, 8), (8, 12), (8,)]
n = sum(torch.Size(s).numel() for s in shapes)
total = 0
for s in shapes:
    w16 = torch.zeros(s, dtype=torch.bfloat16)
    g16 = torch.zeros(s, dtype=torch.bfloat16)
    master = torch.zeros(s, dtype=torch.float32)
    m = torch.zeros(s, dtype=torch.float32)
    v = torch.zeros(s, dtype=torch.float32)
    total += sum(t.numel() * t.element_size() for t in (w16, g16, master, m, v))
assert total == static_train_bytes(n) == 16 * n

psi = 6_738_415_616
print(static_train_bytes(psi) / 1e9, "GB")     # 107.81 GB
print(static_train_bytes(psi) / 2**30, "GiB")  # 100.41 GiB
# pure FP32 training: w4 + g4 + m4 + v4, also 16 bytes
assert static_train_bytes(psi, w=4, g=4, master=0) == static_train_bytes(psi)

The last line is easy to miss: both selected configurations use 16Ψ. This depends on the dtypes of weights, gradients, master weights and optimizer states; other mixed-precision configurations differ. Activation memory and compute costs are separate. The static part of a 7B model alone exceeds an 80 GB card, which is why ZeRO shards optimizer states.

Step 3: FLOPs: 2mnk for matmul, 6ND for training, 8Nd²+4N²d for attention

What's new: check hand-counted FLOPs against torch.utils.flop_counter.FlopCounterMode . (Pro lesson)

Step 4: KV cache bytes and allocatable blocks

What's new: the KV cache byte formula, plus nano-vllm's "carve the free memory into blocks" calculation. (Pro lesson)

Step 5: roofline, why decode is memory bound, TTFT/TPOT

What's new: arithmetic intensity AI = FLOPs / bytes , the roofline time = max(FLOPs/peak, bytes/bw) , and the latency formula. (Pro lesson)

Separate queueing, first token and decode

A request queues, prefills, then decodes. TTFT includes waiting before the first token; TPOT describes later average speed. (Pro lesson)

In Real Systems

nano-vllm nanovllm/engine/model_runner.py , allocate_kv_cache : warm up to measure peak memory, compute block_bytes = 2 * num_hidden_layers… (Pro lesson)

Common Pitfalls

Written: K/V projections counted as d*d . Problem: in GQA models K/V output H_kv*hd ; (Pro lesson)

Say It in the Interview

"For LLaMA-2-7B in FP16, the KV cache is 2 × 32 layers × 32 KV heads × 128 × 2 bytes = 512 KiB per token, so a 4K context costs 2 GiB; (Pro lesson)

Hands-on Notebook

Each exercise in the notebook has a problem statement and a function skeleton; write your version, then run the test cell below it. (Pro lesson)