Non-Weight Memory: Activation Memory, CUDA Context, and System Overhead

When deploying large generative models, software engineers frequently assume that once the model weights fit inside the GPU's High Bandwidth Memory (HBM), the serving system is safe from Out-of-Memory (OOM) crashes.

In production, over 70% of unexpected inference crashes are caused by Non-Weight Memory spikes—specifically transient Activation Tensors, CUDA Allocator Memory Fragmentation, and KV Cache exhaustion.

Understanding how memory dynamically expands and contracts during different execution phases is mandatory for building resilient, high-concurrency LLM inference gateways.

---

1. The Anatomy of Activation Memory

Activation Memory refers to the intermediate tensor representations created in GPU VRAM during the forward pass of a transformer layer.

During training, all activation tensors must be stored in VRAM to compute gradients during backpropagation. During inference, activation tensors only need to exist temporarily while their downstream layer is computing. However, during the Prompt Prefill phase (when thousands of input tokens are processed simultaneously), intermediate activations spike dramatically.

Where Activations Occur in a Transformer Block

In a single transformer decoder layer:

1. Layer Normalization (RMSNorm): Input tensor shape $[B,S,d_{\text{model}}]$.
2. Q, K, V Linear Projections: Generates 3 intermediate matrices of shape $[B,S,d_{\text{model}}]$.
3. Self-Attention Score Matrix:
   • In Standard Attention: Requires computing the full attention score matrix $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^T$, having shape $[B,n_{\text{heads}},S,S]$.
   • In FlashAttention: Avoids materializing the full $[S,S]$ matrix by computing tiled blocks in on-chip SRAM.
4. Feed-Forward Network (SwiGLU): Generates 2 intermediate gate and up projections of shape $[B,S,d_{\text{ffn}}]$ (where $d_{\text{ffn}}\approx \frac{8}{3}{d}_{\text{model}}$).
5. Output Projection & Residual Addition: Shape $[B,S,d_{\text{model}}]$.

---

2. Quadratic vs Linear Scaling: Standard Attention vs FlashAttention

The choice of attention kernel dictates whether activation memory scales quadratically ($\mathcal{O}(S^2)$) or linearly ($\mathcal{O}(S)$) with sequence length.

Standard PyTorch Attention

In standard PyTorch attention, the raw attention matrix $\mathbf{S}\in {\mathbb{R}}^{B\times n_{\text{heads}}\times S\times S}$ is fully written to GPU HBM before Softmax is applied.

${\text{Memory}}_{\text{Attn Matrix}}=B\times n_{\text{heads}}\times S^2\times 2\text{ bytes (in FP16)}$

Scaling Analysis for a 70B Model ($n_{\text{heads}}=64,B=1$):

• At $S=2,048$ tokens: $\text{Memory}=1\times 64\times (2048{)}^2\times 2=0.53\text{ GB}$
• At $S=8,192$ tokens: $\text{Memory}=1\times 64\times (8192{)}^2\times 2=8.58\text{ GB}$
• At $S=32,768$ tokens: $\text{Memory}=1\times 64\times (32768{)}^2\times 2=137.4\text{ GB}$ (Guaranteed OOM!)

FlashAttention-2 / FlashAttention-3

FlashAttention computes attention in tiled blocks inside the Streaming Multiprocessor's fast on-chip SRAM ($228\text{ KB}$ per SM), maintaining running Softmax reduction statistics without ever materializing the $S\times S$ matrix in VRAM:

${\text{Peak Activation Memory}}_{\text{FlashAttention}}=\mathcal{O}(B\times S\times d_{\text{model}})$

• At $S=32,768$ tokens on Llama 3 70B ($d_{\text{model}}=8192$): $\text{Memory}\approx 1\times 32768\times 8192\times 2\text{ bytes}\approx 0.53\text{ GB}$

---

3. PyTorch Caching Allocator and Memory Fragmentation

A major hidden source of non-weight VRAM overhead is CUDA Memory Fragmentation.

When PyTorch executes neural network kernels, it does not call the low-level OS cudaMalloc and cudaFree for every tensor because dynamic allocation in the GPU driver is extremely slow (~10–20 microseconds per call). Instead, PyTorch uses a Caching Allocator.

Internal vs External Fragmentation

1. External Fragmentation: Free memory is split into numerous non-contiguous small memory chunks. When a large activation tensor requires a contiguous block (e.g. 2 GB), cudaMalloc fails even though the sum of free fragments exceeds 4 GB.
2. Reserved vs Allocated Memory:
   • torch.cuda.memory_allocated(): Memory currently occupied by live tensors.
   • torch.cuda.memory_reserved(): Total memory held by the caching allocator from the GPU driver.
   • The gap between reserved and allocated is cached, idle memory that is unavailable to other processes.

---

4. Measuring Real-Time Non-Weight Memory Dynamics

Below is an interactive Python utility demonstrating how to track live activation surges, measure allocator fragmentation, and tune allocator settings.

---

5. Production Failure Modes and Engineering Runbook

Failure Mode 1: Activation OOM During Long Document Analysis

• Symptom: A chatbot works flawlessly on short prompts (500 tokens), but instantly crashes with CUDA OOM when a user uploads a 50-page PDF (25k tokens).
• Root Cause: The inference engine is using a naive PyTorch attention implementation instead of FlashAttention. When $S$ expanded from $500$ to $25,000$, activation memory exploded by $50^2=2,500\times$, requiring over 60 GB of intermediate attention matrix memory.
• Resolution: Verify and enforce FlashAttention-2 / vLLM kernel integration (model.config._attn_implementation = "flash_attention_2"). Ensure the FlashAttention CUDA binaries are compiled for the host GPU architecture.

Failure Mode 2: PyTorch Allocator Fragmentation OOM

• Symptom: torch.cuda.OutOfMemoryError: Tried to allocate 512 MB. GPU total: 80 GB, Available: 14 GB.
• Root Cause: External memory fragmentation prevents the allocator from finding a contiguous 512 MB chunk, despite having 14 GB of total free fragmented space.
• Resolution: Configure the PyTorch allocator with the max_split_size_mb environment variable:

This instructs the allocator to prevent splitting large memory blocks into tiny unusable fragments.

---

6. Summary & Key Takeaways

1. Prefill Triggers Activation Bursts: The prompt prefill phase creates large intermediate tensors across feed-forward and attention layers that peak significantly above baseline weight memory.
2. FlashAttention Eliminates Quadratic Scaling: By tiling attention computation inside on-chip SRAM, FlashAttention reduces peak activation VRAM from $\mathcal{O}(S^2)$ to $\mathcal{O}(S)$.
3. Beware Allocator Fragmentation: PyTorch's caching allocator holds reserved memory that can become fragmented, causing OOMs even when total available VRAM appears sufficient.
4. Tune Allocator Guardrails: Set PYTORCH_CUDA_ALLOC_CONF and reserve a 10–15% safety margin in serving frameworks to handle unpredictable prompt lengths.