GPU Memory Management for LLM Inference: Beyond Model Weights

Self-hosting large language models (LLMs) has become a rite of passage for modern developers. Whether you are building a private RAG (Retrieval-Augmented Generation) pipeline or integrating LLMs into a LangChain workflow, the first hurdle is always hardware. However, most 'run an LLM on your GPU' tutorials suffer from a fatal flaw: they focus exclusively on model weights. They show you how to load a 7B model, run a simple 'Hello World' prompt, and declare victory.

Then, the moment you send a 6,000-token document at a batch size of eight, the system crashes with a dreaded Out-of-Memory (OOM) error. While platforms like n1n.ai provide a seamless way to bypass these hardware headaches with stable, high-speed APIs, understanding the underlying arithmetic is crucial if you insist on managing your own infrastructure. The failure isn't random—it's predictable physics.

The Four Buckets of GPU Memory

GPU memory for inference is not a monolithic block. It is divided into four separate buckets, and tutorials usually only mention the first one. To successfully deploy models like Llama-3.1-8B or even DeepSeek-V3, you must account for all of them.

1. Model Weights: This is the static portion. It is fixed, predictable, and loaded once into VRAM.
2. KV Cache: This is the per-request memory that stores context. It scales linearly with context length and batch size. This is what usually kills your process.
3. Activations and CUDA Overhead: The runtime itself (PyTorch, CUDA context) plus temporary tensors generated during a forward pass.
4. Fragmentation: Memory that is technically free but unusable because it is scattered in non-contiguous blocks.

Bucket 1: The Weights (The Easy Part)

The weights number is what everyone quotes because the math is simple. You take the parameter count and multiply it by the bytes-per-parameter based on your precision (FP16, INT8, or INT4).

If you have an RTX 3090 or 4090 with 24GB of VRAM, a 4-bit 7B model looks like it leaves 20GB of 'spare' space. That space is not spare; it is your working budget for the next three buckets. If you find this overhead too complex to manage for production, using a managed provider like n1n.ai allows you to focus on the application logic rather than VRAM arithmetic.

The Hidden Killer: KV Cache Arithmetic

The KV (Key-Value) cache stores the attention mechanism's tensors for every token already processed. This prevents the model from recomputing the entire history at every new token generation step. The size of this cache per request is roughly:

kv_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element

Modern models like Llama-3 or DeepSeek-V3 use Grouped-Query Attention (GQA), which significantly reduces the num_kv_heads compared to standard multi-head attention. This is a massive win for memory, but the cache still grows fast.

Worked Example: Llama-3-8B (FP16 KV Cache)

• Layers: 32
• KV Heads: 8
• Head Dim: 128
• Precision: 2 bytes (FP16)

Per-token memory = 2 * 32 * 8 * 128 * 2 = 131,072 bytes (approx. 128 KB/token)

At a context length of 8,192 tokens, a single request consumes 1 GB of VRAM just for the KV cache. If you want to serve a batch of 16 users simultaneously, that is 16 GB of VRAM. Suddenly, that 'extra' 20GB on your RTX 4090 is almost entirely gone before you even account for CUDA overhead.

Bucket 3 & 4: Overhead and Fragmentation

Even with zero tokens, loading the CUDA context and your inference framework (like PyTorch) claims a 'tax.' This is usually between 1GB and 2GB. Frameworks like vLLM are proactive; they deliberately grab up to 90% of available VRAM (controlled by the gpu_memory_utilization flag) to manage the KV cache internally.

Fragmentation is the second silent killer. In a naive implementation, if you process requests of varying lengths, the memory becomes 'Swiss cheese.' You might have 4GB of total free space, but no single contiguous 2GB block. This leads to an OOM even when nvidia-smi says you have space.

This is why PagedAttention (pioneered by vLLM) is so critical. It treats VRAM like virtual memory in an operating system, breaking the KV cache into non-contiguous blocks (pages). If you are moving from local development to production, you should prioritize frameworks that support PagedAttention.

Choosing Your Stack: A Comparison

For most developers starting out, Ollama is the gold standard for ease of use. However, as you scale toward supporting multiple users or long-context RAG applications, the limitations of simple wrappers become apparent. At that stage, you either need to master vLLM configuration or migrate to a robust API aggregator like n1n.ai to handle the infrastructure scaling for you.

The Pre-Flight Check: VRAM Estimation Script

Before you rent an H100 or buy a 4090, run this simple Python estimation. It will tell you if your target concurrency is realistic.

def estimate_vram(params_b, bits, layers, kv_heads, head_dim, ctx, batch):
    # Weight memory
    weight_gb = (params_b * (bits / 8))

    # KV Cache per token in GB
    # 2 for Key and Value
    bytes_per_param = 2 # Assuming FP16 KV cache
    per_token_gb = (2 * layers * kv_heads * head_dim * bytes_per_param) / 1e9

    total_kv_gb = per_token_gb * ctx * batch
    overhead_gb = 2.0 # CUDA + Framework base

    return round(weight_gb + total_kv_gb + overhead_gb, 2)

# Example: Llama-3-8B, 4-bit, 16k context, batch of 4
print(f"Estimated VRAM: {estimate_vram(8, 4, 32, 8, 128, 16384, 4)} GB")

Pro-Tips for Optimization

1. KV Cache Quantization: You can store the KV cache in FP8 or even INT8. This halves the memory requirement with negligible impact on perplexity. vLLM and TGI both support this, though it is often off by default.
2. Context Capping: Do not set your max_model_len to 128k just because the model supports it. Set it to the actual maximum your application needs. Many frameworks pre-allocate based on this ceiling.
3. GQA Awareness: When picking a model, check if it uses Grouped-Query Attention. Models like Mistral and Llama-3 use it, making them much more memory-efficient than older models like GPT-NeoX.

Conclusion

Self-hosting an LLM is a balance of weights and cache. The weights determine if the model starts; the KV cache determines if the model serves. By doing the math upfront, you can avoid costly OOM errors and hardware misinvestments.

If the complexity of managing CUDA versions, PagedAttention blocks, and VRAM fragmentation is taking too much time away from your core product development, consider using n1n.ai. We aggregate the world's most powerful models, including Claude 3.5 Sonnet and OpenAI o3, into a single, high-performance API so you can scale without worrying about the hardware.

Get a free API key at n1n.ai