This is my quick access note for kv, flops, roofline analysis for llm inference.

KV size

This is a warm up practice to massage your brain. Let’s calculate the KV size of a single token for llama3 8B model.

Llama3 8B model with sequence length 1024

• seq_len = 1024
• d_model = 4096
• n_q_head = 32
• n_kv_head = 8 → GQA factor = 4
• d_head = 128 (hidden dimension per head)
• d_type = 2 (fp16)
• d_ffn = 14336
• num_layers = 32
• vocab_size = 128256

KV size per token = num_layers * n_kv_head * d_head * d_type * 2(K and V) 
                  = 32 * 1024 * 2 * 2
                  = 131,072 bytes (0.125 MB)

for 1024 tokens, the total KV size is 1024 * 0.125 MB = 128 MB

KV read latency

This table compares KV cache theorectical read latency from different memory hierarchy. The data size represents the total KV cache footprint for 1024 tokens for llama3 8B model. It does not have analysis when kv cache is on CXL pool.

Typical NVMe data path is Disk → PCIe → CPU(DRAM) → PCIe → GPU.

• Bandwidth: ~5 GB/s

NVMe GPUDirect data path: Disk→PCIe→GPU (direct DMA from NVMe to GPU memory, single hop, bypasses CPU)

• Bandwidth: ~7 GB/s

FLOPs table for different models

Q projection: 2×seq×d_m×d_q

K projection: 2×seq×d_m×d_kv

V projection: 2×seq×d_m×d_kv

QxK^T: 2×seq²×d_q

Attn×V: 2×seq²×d_q

O projection: 2×seq×d_q×d_m

MLP up: 2×seq×d_m×d_ff

MLP down: 2×seq×d_ff×d_m