Skip to content

Research 02 — Analysis engine math: params, FLOPs, KV cache, memory, cost

Researched 2026-09-18. Formulas were validated numerically against published configs (see §4 and params_reference.py).

Notation: L layers, D d_model, H query heads, Hkv KV heads, g = H/Hkv, dh head dim, F FFN width, V vocab, T sequence length, B batch, E routed experts, k top-k, n_s shared experts, Fe expert FFN width, N parameters, bytes = dtype width (bf16 = 2).

1. Formulas and references

Canonical sources: EleutherAI Transformer Math 101; Kaplan scaling laws; Chinchilla; PaLM App. B for the 6N + 12LHQT rule; Megatron activation recomputation, Korthikanti et al.; ZeRO; FlashAttention; GQA; MQA; RoPE; DeepSeek-V2 for MLA; DeepSeek-V3; DeepSeekMoE; Mamba-2 / SSD and reference code; kipply Transformer Inference Arithmetic for roofline; 8-bit Adam.

1.1 Parameter count (validated, see §4)

Block Params
Token embedding V·D
LM head V·D if untied, 0 if tied (GPT-2, Gemma tie; Llama/Qwen/DeepSeek don't)
Learned positions (GPT-2) T_max·D; RoPE/NoPE: 0
Attention MHA/GQA/MQA D·H·dh + 2·D·Hkv·dh + H·dh·D (+ biases H·dh + 2·Hkv·dh + D if attention_bias; Qwen2.5 has QKV bias, no O bias). MHA: Hkv=H; MQA: Hkv=1. QK-norm (Qwen3, Gemma 3): +2·dh
Attention MLA (DeepSeek-V2/V3) Q path: D·d_c' + d_c' + d_c'·H·(d_nope + d_rope); KV path: D·(d_c + d_rope) + d_c + d_c·H·(d_nope + d_v); O: H·d_v·D. V3 (D=7168, H=128, d_c'=1536, d_c=512, d_nope=128, d_rope=64, d_v=128) → 187.1 M/layer
Dense MLP, 2-matrix (GELU/ReLU²) 2·D·F (+ F + D biases)
Gated MLP (SwiGLU/GeGLU) 3·D·F (Llama: F = round_to_256(8/3·D) → 11008, 14336)
MoE layer router D·E (+E if bias, e.g. DeepSeek e_score_correction_bias) + routed E·3·D·Fe + shared n_s·3·D·Fe. Respect dense-layer overrides (first_k_dense_replace, decoder_sparse_step, mlp_only_layers)
Norms RMSNorm D each; 2/layer pre-norm, 4/layer for Gemma (pre+post); +1 final. LayerNorm 2D
Mamba-2 block (d_inner=E·D, nheads=d_inner/P, groups G, state N_s) in_proj D·(2·d_inner + 2·G·N_s + nheads) + conv1d (d_inner + 2·G·N_s)·(d_conv+1) + 3·nheads (dt_bias, A_log, D) + norm d_inner + out_proj d_inner·D. Defaults N_s=128, P=64, E=2, d_conv=4
Active params (MoE) same, but k·3·D·Fe + n_s·3·D·Fe instead of E·…; embeddings/head count once

1.2 FLOPs per token

  • Rule of thumb: forward ≈ 2·N (N = non-embedding active params, LM head included since it is a matmul; embedding lookup is a gather ≈ 0). Backward ≈ 4·N. Train 6·N; with full activation recompute 8·N.
  • Attention (score + PV) term, per token, per layer, forward: 4·T·H·dh (QKᵀ 2·T·H·dh, PV 2·T·H·dh), independent of Hkv. Halve if the kernel skips masked causal blocks (FlashAttention does). Full training: FLOPs/token = 6·N + 12·L·T·H·dh (PaLM App. B form).
  • Where 6N/2N break: (a) long context: with N ≈ 12·L·D² the attention share is T/(6·D) → 33% for Llama-3-8B at 8K, >100% at 32K; (b) MoE: use active N, but hardware FLOPs include padded/dropped tokens under capacity factors; (c) attention is 2·T·H·(d_qk + d_v) per token before the causal half, the scores at the key width and the value product at the value width — they differ in latent attention (192 against 128 in DeepSeek-V3) and in differential attention (values twice the key width); MLA decode with weight absorption scores in latent space: 2·T·H·(d_c + d_rope) per token, larger than GQA's 2·T·H·dh; (d) MTP adds ≈ one extra layer + head per depth; (e) sliding-window layers: a causal window of W keeps min(q + 1, W) keys at position q, so replace the causal T/2 by W − W²/2T once T > W (not W/2: the window already only looks back); a mask written as an expression is counted as the share of those it keeps, measured; sparse attention by topk + W; (f) SSM layers: linear in T; SSD scan ≈ 6·d_inner·N_s FLOPs per token plus a chunk-quadratic term ≈ 4·d_inner·Q for chunk Q=256, small vs the in/out_proj matmuls (approximate; paper gives asymptotics not constants); (g) elementwise ops (norm, RoPE, softmax, act) are ~1–3% of FLOPs but memory-bound: ignore for FLOPs, not for latency; (h) inference decode per token: 2·N_active + 4·L·T_ctx·H·dh.

1.3 KV-cache bytes per token

  • MHA/GQA/MQA: 2·L·Hkv·dh·bytes. Llama-2-7B (MHA) 512 KiB; Llama-3-8B (GQA g=4) 128 KiB; Llama-3-70B 320 KiB (the 160 KiB first reported here was wrong: 80 layers x 8 kv heads x 128 head dim x 2 tensors x 2 bytes = 320 KiB, confirmed by the implementation's regression test). DeepSeek-V4 uses Hkv=1 with head_dim=512 and K=V shared → L·512·bytes.
  • MLA: L·(d_c + d_rope)·bytes (DeepSeek-V2 Table 1). V3: 61·576·2 = 70,272 B ≈ 68.6 KiB/token vs 4 MiB uncompressed (57×). Engines that don't absorb weights cache decompressed K/V → MHA size; the analysis exposes both, as bytesPerToken and bytesPerTokenDecompressed. The decompressed figure is a key and a value per head plus the rotary part, which every head shares and so is held once: L·(H·(d_nope + d_v) + d_rope)·bytes. For V3 that is 3.82 MiB against 68.63 KiB, reproducing Table 1's 57x, and a test pins the ratio. A design with no latent attention reports the same number twice rather than a second figure that means nothing.
  • Sliding-window layers: cache min(T, W) tokens, so per-sequence cost 2·L_local·Hkv·dh·bytes·W is constant. Gemma-3-27B (5:1 local:global, W=1024, 62 layers): ~11 global layers → 90 KB/token, plus ≈ 428 MB fixed for 51 local layers.
  • SSM layers: fixed per sequence, not per token: nheads·P·N_s·bytes + (d_inner + 2·G·N_s)·(d_conv−1)·bytes. Nemotron-H-8B: ≈ 50 MB/sequence for 24 Mamba layers, plus 4 attention layers × 16 KiB/token.
  • Serving total: per_token × T_ctx × concurrent sequences (+ paged-attention block waste).

1.4 Training memory

  • Per parameter (mixed precision AdamW): bf16 weights 2 + bf16 grads 2 + fp32 master 4 + m 4 + v 4 = 16 B/param. fp32 grads → 18 B. 8-bit Adam → 10 B. SGD+momentum 8 B optimizer; Muon ≈ 8 B. Pure-bf16 Adam without master weights: 8 B (unstable; flag in UI).
  • Activations per layer (Megatron eq. 1, GPT-3-style block, bf16, F=4D, dropout): s·b·h·(34 + 5·a·s/h) bytes, i.e. per token 34·D + 5·H·T. With tensor parallel t: sbh(10 + 24/t + 5as/(ht)); with sequence parallel: sbh(34/t + 5as/(ht)); selective recompute or FlashAttention drops the 5as/h term → 34·sbh/t; full recompute 2·sbh (layer input only); recompute costs +33% FLOPs (8N).
  • Modern SwiGLU/no-dropout block (own derivation from saved tensors; verify against a profiler): per token per layer ≈ D·(14 + 4/g) + 6·F bytes bf16 (norm inputs 4D, QKV input 2D, Q/K/V for FA backward 2D(1+2/g), FA output + o_proj input 4D, MLP input 2D, gate/up outputs 4F, act output 2F). Llama-3-8B: ≈ 147 KB/token/layer → 38 GB for one 8K sequence without recompute. Implementation note: this list omits the down-projection's input. Unfused PyTorch keeps four F-wide tensors (gate output, activation output, up output, product), not three, so the real figure is ≈ 176 KB/token/layer. A fused gated feed-forward reaches the 147 KB quoted here.
  • Logits: B·T·V·(2 + 4) bytes (bf16 logits + fp32 softmax/CE); 8K × 128K vocab = 6.3 GB. Use chunked CE.
  • Measured (tensorcad-runtime measure on an RTX 5080, PyTorch 2.11, B=2, T=512):
  • At rest. Weights, gradients and optimizer state agree with the 16 B/param above to within 2%. PyTorch's mixed precision reaches the same sixteen bytes by another route: fp32 weights, fp32 gradients and two fp32 moments.
  • Saved activations with the whole model in bf16. The activation dtype assumed here, and within 6% of the analysis: GPT-2 small 588 MiB measured against 604 analysed, a Llama-style 150M model 896 against 848.

  • Under autocast. Plain PyTorch training saves 1.26–1.55× that: 933 and 1,071 MiB. precision: autocast models it, and comes to 894 and 1,018 MiB (0.96× and 0.95× of the measurement), and 1,552 MiB against 1,586 for GPT-2 at T=1024. The model has three parts, in analysis/autocast.go:

    • Width. Each tensor's width follows the graph. An embedding lookup is fp32, since the table is. Norms, softmax, rotations and the scans are fp32, by autocast's own list and by what they are written in. Matrix multiplies and the attention kernel produce bf16. Anything else is the widest of what it reads, so the residual stream is fp32 from the embedding on.
    • Per-consumer casts. A matrix multiply that reads an fp32 tensor saves its own bf16 cast of it, one per reader. Three projections of one norm output are three copies.
    • Weight copies. Every weight a matrix multiply casts is held as a bf16 copy until the backward pass has used it: 2 B per such parameter, a tied head's table included.

    At rest the split is 4 + 4 + 8 B/param (fp32 weights are their own master copy), the same 16 as mixed precision, and it matches the allocator to 1%. - Sharding: ZeRO-1 W + G + Opt/N_dp + Act; ZeRO-2 W + (G+Opt)/N_dp + Act; ZeRO-3/FSDP (W+G+Opt)/N_dp + Act + live unsharded layer. TP divides W/G/Opt by t and activations by t only with sequence parallel; PP divides layers by p but holds p micro-batches of activations in flight (1F1B); EP divides only expert weights. Add 10–20% workspace/fragmentation.

1.5 Inference memory and roofline

  • Memory = N_total·bytes (all experts resident) + KV cache + small activations + ~20% overhead.

  • Decode step time ≈ max(bytes_moved / BW, FLOPs / (peak·eff)); bytes_moved = read weight bytes + KV bytes(B·T); FLOPs = B·(2·N_active + 4·L·T·H·dh). Arithmetic intensity of weight streaming ≈ B FLOP/byte (bf16), so decode is compute-bound only when B > ridge = peak/BW: A100 ≈ 208; H100 SXM ≈ 989/3.35 ≈ 295; B200 ≈ 2250/8 ≈ 280 (approximate spec values). Batch-1 tok/s upper bound = BW / weight_bytes (Llama-3-8B bf16 on H100 ≈ 209 tok/s). Prefill is compute-bound: 2·N·T/(peak·MFU).

  • A sparse model at batch reads more than one token's share. One token routes to k of E experts and reads k/E of them, which is the active count. A batch of B tokens routes independently, so an expert is read unless every one of them skipped it: the expected share read is 1 - (1 - k/E)^B, climbing from k/E at B=1 to essentially 1 well before B reaches E. Mixtral at batch 32 reads 43.49 GiB of the 43.50 GiB it holds, not the 12.9 GiB one token reads, and its decode rate is 929 tok/s where the active count alone says 1314. The implementation writes this per block as Σ per · copies · (1 - (1 - active/copies)^B), which needs no knowledge of k or E and reduces to the active count at B=1 and to the resident count for anything dense. This assumes routing is independent and uniform; a router trained towards balance spreads a batch over more experts still, so the figure is a floor on what is read and the rate above it a ceiling.

1.6 Training cost

FLOPs_total = (6·N_active + 12·L·T·H·dh)·tokens (×8/6 with full recompute); GPU-hours = FLOPs_total / (peak_FLOPs · MFU · 3600); wall-clock = GPU-hours / n_GPU; cost = GPU-hours · $/h.

MFU anchors: Llama-3-405B 43%/41%/38% BF16 on 8K/16K/16K+CP H100s (Llama 3 paper); PaLM 46.2% on TPUv4; DeepSeek-V3: 2.664 M H800-hours for 14.8 T tokens → ≈ 38% BF16-equivalent (~19% of FP8 peak). Suggested UI ranges (synthesis): dense well-tuned 35–45%; MoE 20–35%; naive/small clusters 25–35%.

1.7 Chinchilla and over-training

L(N,D) = E + A/N^α + B/D^β; Hoffmann: E=1.69, A=406.4, B=410.7, α=0.34, β=0.28; Epoch refit E=1.8172, A=482.01, B=2085.43, α=0.3478, β=0.3658. Compute-optimal D_opt ≈ 20·N, N_opt = sqrt(C/120). Real tokens/param ratios (Epoch): Llama-2-7B 286; Llama-3-70B 214; Llama-3.1-405B 38.5; Llama-3-8B ≈ 1875; Qwen2.5-72B 248; DeepSeek-V3 22 (total) / 400 (active); Qwen3-0.6B ≈ 60,000. Trend: 3.1×/yr. Expose both total- and active-param ratios; cite inference-adjusted optima (Sardana & Frankle).

2. Existing calculators / libraries

Tool Covers Gaps
EleutherAI cookbook calc/ params (GQA, MoE, tied emb), FLOPs (train/checkpointing, attention terms, MoE gating), memory (ZeRO 1-3, TP/PP, Adam/8-bit/SGD/Muon, Megatron activation formulas, KV cache) no MLA, no SSM, no sliding window; KV formula may omit kv-ratio (verify)
llm-analysis (Apache-2) train+infer latency/memory/FLOPs, TP/PP/SP/EP/DP, dtype bits, efficiency knobs GQA/MLA/SSM/FlashAttention not modelled
LLM-para (MIT) inference roofline over 13 operators; GQA, MoE, MLA, RoPE, SwiGLU, FA; 24 hardware targets inference only; no SSM
torch.utils.flop_counter.FlopCounterMode dispatch-level, shape-based; mm/bmm/conv/SDPA fwd+bwd; custom ops via register_flop_formula; works with fake/meta tensors only registered aten ops. Good validation backend.
DeepSpeed FLOPS profiler per-module breakdown custom kernels unseen
calflops HF models; FLOPs/MACs/params same limits
torchinfo, fvcore params, mult-adds by hooks miss functional matmuls, no training memory
HF Model Memory Utility weight bytes, training ≈ 4× no KV/activations, no FLOPs
Vokturz can-it-run-llm, apxml, NyxKrage, modelfit.io inference/fine-tune memory vs GPU opaque or HF-only

Nothing covers MoE + MLA + SSM + training + inference together; SSM/hybrid accounting is the clearest gap.

3. Shape inference / symbolic approach

  • Representation: each port carries (dims: SymExpr[], dtype) where SymExpr is a canonical integer polynomial over named symbols (B, T, D, H, dh, F, V…) with rational coefficients. Covers concat (+), split/reshape (*, / with a divisibility obligation like D = H·dh), and constants. Validation = unify producer shape with consumer pattern: bind free names, compare canonical forms, emit diffs ("expected (B,T,D), got (B,T,H·dh) and H·dh ≠ D"). Keep a global constraint set (F % 256 == 0, H % Hkv == 0, E ≥ k) checked as design rules.
  • Port-type notation: jaxtyping strings ("batch seq dim", expression axes, variadic "*batch") and einops patterns ("b t (h d) -> b h t d") are a ready-made port-typing DSL. Adopt that syntax.
  • Reference implementations: torch.export uses SymPy-backed SymInts in a ShapeEnv with guards, Dim(name, min, max) (docs); ONNX dim_param strings with no arithmetic; onnxruntime SymbolicShapeInference uses SymPy.
  • JS/TS: Algebrite, math.js, nerdamer are heavier than needed; a ~300-line polynomial canonicalizer is the pragmatic choice.
  • Validation backend: code-gen PyTorch from the graph, run torch.export with Dim("B"), Dim("T") to cross-check shapes, and FlopCounterMode on meta/fake tensors to cross-check FLOPs (meta-device counting should work since counting is shape-based; verify).

4. Known-good architecture specs (regression set)

Configs from HF config.json or the paper; "calc" = params_reference.py in this folder.

Model L D H / Hkv / dh F (Fe) V Experts / k / shared Published Calc
GPT-2 small/med/large/XL (tied, learned pos, biases, LN) 12/24/36/48 768/1024/1280/1600 12/16/20/25, MHA 4D 50257 – 124M/355M/774M/1.5B 124.4/354.8/774.0/1557.6M
Llama-2-7B 32 4096 32/32/128 11008 32000 – 6.74B 6.738B
Mistral-7B 32 4096 32/8/128 14336 32000 – 7.24B 7.242B
Llama-3-8B 32 4096 32/8/128 14336 128256 – 8.03B 8.030B
Llama-3-70B 80 8192 64/8/128 28672 128256 – 70.6B 70.55B
Llama-3.1-405B 126 16384 128/8/128 53248 128256 – 405B 405.9B
Mixtral-8x7B 32 4096 32/8/128 (14336) 32000 8 / 2 / 0 46.7B / 12.9B act 46.70B / 12.88B
Qwen2.5-7B (QKV bias) 28 3584 28/4/128 18944 152064 – 7.61B 7.615B
Qwen3-8B (QK-norm) 36 4096 32/8/128 12288 151936 – 8.2B 8.19B
Qwen3-30B-A3B 48 2048 32/4/128 (768) 151936 128 / 8 / 0 30.5B / 3.3B 30.53B / 3.35B
Qwen3-235B-A22B 94 4096 64/4/128 (1536) 151936 128 / 8 / 0 235B / 22B 235.1B / 22.2B
DeepSeek-V3 MLA d_c'=1536, d_c=512, nope 128, rope 64, d_v 128; 3 dense layers F=18432; router bias 61 7168 128 MLA (2048) 129280 256 / 8 / 1 671B / 37B; +14B MTP 671.0B / 37.6B; MTP 13.5B
Gemma-2-9B (tied, 4 norms/layer, GeGLU) 42 3584 16/8/256 14336 256000 – 9.24B 9.242B
Gemma-3-27B (tied; W=1024, 5:1) 62 5376 32/16/128 21504 262208 – text 27.0B 27.01B
Nemotron-H-8B: 24×Mamba-2, 4×attn, 24×FFN(relu², 2-matrix); Mamba-2 nheads 128, P 64, N_s 128, G 8 52 4096 32/8/128 21504 131072 – ~8B 8.10B
Jamba-v0.1: attn every 8, MoE every 2, Mamba-1 32 4096 32/8/128 14336 65536 16 / 2 52B / 12B not computed

2026 models (from HF docs/config and secondary sources; totals not recomputed): DeepSeek-V4-Flash 284B/13B — 43 layers, D=4096, 256+1 experts, top-6, Fe=2048, shared-K=V MQA with head_dim 512, grouped low-rank O-proj, sliding window 128 + compressors (HF docs). Kimi-K3 — 93 layers, D=7168, 96-head MLA, 896 routed + 2 shared, top-16, KDA linear-attention 3:1, V=163840; reported 2.8T/104B. GLM-5 744B/40B — 78 layers, D=6144, 256 experts top-8 + 1 shared, MLA + DSA. Qwen3.5-397B-A17B and Qwen3-Coder-Next 80B-A3B — Gated DeltaNet + gated attention 3:1 (Raschka).

Flags / unverified

  • Mamba-2 SSD FLOP constants (§1.2 f) and the SwiGLU activation-bytes formula (§1.4) are derivations, not quoted from papers.
  • H100/B200 ridge points use approximate spec-sheet values.
  • 2026 model totals (Kimi-K3, GLM-5, DeepSeek-V4) come from vendor/secondary summaries; recompute before using as tests.
  • Jamba was not recomputed.
  • Ultra-Scale Playbook is a useful memory/parallelism reference but could not be fetched (JS-rendered).