KV Cache VRAM Explosion When Deploying Multimodal LLMs
KV Cache VRAM Explosion When Deploying Multimodal LLMs
Search
Ask the AI

KV Cache VRAM Explosion When Deploying Multimodal LLMs

Introduction: The Deployment Dilemma of MLLMs

When deploying Multimodal Large Language Models (MLLMs) such as LLaVA, Qwen-VL, or InternVL, developers frequently encounter a critical memory bottleneck: the KV Cache VRAM explosion. Particularly when processing high-resolution (e.g., 4K) images or lengthy videos, the number of visual tokens scales with pixel area, and the sequence rapidly outgrows both the context window and the prefill compute budget. This post gets the KV cache arithmetic right first, then works through the order in which optimisations are worth applying.

Mathematical Derivation of KV Cache Memory Consumption

During the autoregressive decoding phase of the Transformer architecture, we cache the Keys and Values of previous tokens in VRAM to prevent redundant computations—this is known as the KV Cache. The memory footprint for the KV Cache of a single token can be calculated precisely as follows:

Formula: $VRAM_{KV} = 2 \times layers \times kv\_heads \times head\_dim \times seq\_len \times bytes$

Parameter breakdown:

  • 2: Accounts for the Key and Value tensors.
  • layers: Number of decoder layers in the model (e.g., 32 layers for Llama-3-8B).
  • kv_heads: The number of KV heads, not attention heads. See below.
  • head_dim: Dimension of each head.
  • seq_len: The current sequence length, i.e., the total number of tokens.
  • bytes: Precision format size (e.g., 2 bytes for FP16/BF16, 1 byte for INT8).

Use KV heads, not attention heads

This is where the formula is most often miscomputed. Modern models use grouped-query attention (GQA), where several query heads share one set of Key/Value heads. Llama-3-8B has 32 attention heads but only 8 KV heads. Substituting the attention head count overstates the result by 4×.

seq, layers, head_dim, byt = 8000, 32, 128, 2
kv = lambda h: 2 * layers * h * head_dim * seq * byt / 2**30

kv(32)   # 3.91 GiB  -- attention heads, wrong
kv(8)    # 0.98 GiB  -- KV heads, correct

The corresponding config field is usually num_key_value_heads. Only when a model lacks it (older MHA architectures) does it equal num_attention_heads.

Growth with sequence length is linear, not quadratic

seq_len appears to the first power, so KV cache memory grows linearly with sequence length: double the tokens, double the cache. What grows quadratically is the attention score matrix $QK^T$ — that is compute, not resident memory, and with FlashAttention it is never fully materialised at all.

Separating these two matters, because they call for entirely different remedies: the linear cache is addressed by quantisation and paged management, the quadratic compute by kernel fusion and reducing token count. Conflating them leads to applying the wrong fix.

How many tokens a high-resolution image actually produces:
Assume a ViT-L vision encoder with a patch size of 14×14. A 1024×1024 image is sliced into $(1024/14)^2 \approx 5329$ visual tokens. Processing a 4K image (3840×2160) without any tiling:

(3840 ÷ 14) × (2160 ÷ 14) = 274 × 154 ≈ 42,196 tokens

Note that token count is proportional to pixel area (quadratic in side length). At 42,000 tokens the sequence already exceeds most models’ context windows, and that is the real reason high resolution demands tiling — not that the cache will not fit (roughly 5 GiB by the formula above, comfortable on an A100), but that the sequence does not fit in the context window at all, and prefill attention scales with the square of that length.

AnyRes/Dynamic Resolution: Mitigating Visual Token Explosion

To process high-resolution images without triggering OOM errors, the industry widely adopts the AnyRes (Dynamic Resolution) image tiling strategy. This technique dynamically slices a high-definition image into multiple lower-resolution local patches while preserving an overall global thumbnail.


graph TD
    A[Original High-Res Image 4K] --> B{AnyRes Dynamic Slicing}
    B --> C[Global View 
Resized to 336x336] B --> D[Local Patches
Sliced into NxM 336x336 grids] C --> E(ViT Vision Encoder) D --> E E --> F[Feature Concatenation] F --> G[Projector
MLP Downsampling/Compression] G --> H[LLM Backbone]

This strategy allows the model to capture crisp local details while retaining the overarching semantic context, simultaneously shrinking the sheer volume of visual tokens through pooling or projector layers before they hit the LLM.

Code in Action: 4-bit Quantized LLaVA Inference using llama.cpp

For edge devices or consumer-grade GPUs, we can deploy MLLMs using the C++ backend `llama.cpp` alongside 4-bit quantization (like GGUF) to drastically alleviate VRAM footprint and memory bandwidth bottlenecks.


#!/bin/bash
# Clone and build llama.cpp with CUDA hardware acceleration support
git clone https://github.com/ggerganov/llama.cpp.git
cd llama.cpp
make LLAMA_CUDA=1

# Download quantized multimodal model (LLaVA-1.5-7b) and vision projector (mmproj)
wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/ggml-model-q4_k.gguf
wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/mmproj-model-f16.gguf

# Run llava-cli for multimodal image-text inference
./llava-cli \
  -m ggml-model-q4_k.gguf \
  --mmproj mmproj-model-f16.gguf \
  --image /path/to/high_res_input.jpg \
  -p "Describe the image in detail, paying attention to the intricate textures." \
  -c 4096 \
  -ngl 35 # Offload specific layers to GPU VRAM

The real bottleneck is prefill, not decode

Experience optimising text-only LLMs sends people in the wrong direction here. A text conversation’s prompt runs to a few hundred tokens and nearly all the time goes into decoding token by token, so optimisation focuses on decode-stage memory bandwidth. Multimodal does not work that way.

Visual tokens all live in the prefill — they are part of the prompt and must be processed in one pass before the first word is emitted. That inverts the cost structure:

  • Prefill: n tokens processed at once, with attention compute proportional to $n^2$. At n = 5,000 that is 25 million score computations per head per layer; at n = 42,000 it is 1.76 billion — a 70× increase.
  • Decode: each generated token needs one pass over the KV cache, with compute proportional to $n$ and memory proportional to $n$.

In practice this presents as: ask about a high-resolution image and the first word takes a long time to appear, after which output flows smoothly. The “slowness” a user perceives is almost entirely time-to-first-token, not throughput. If you are optimising output speed and seeing no perceived improvement, you are probably working on the wrong end.

That conclusion sets the optimisation order directly:

  1. Reduce token count first. This is the only lever that improves both ends — prefill benefits quadratically, cache linearly. AnyRes tiling, pooling and Token Merging all belong here. Compressing 42,000 tokens to 3,000 cuts prefill compute to 1/196 of its former cost.
  2. Then compress cache precision. Quantising the KV cache to int8 halves the linear term but does nothing for the quadratic prefill term. It addresses “how many concurrent requests can I serve,” not “how fast does the first token arrive.”
  3. Paged management last. Paging addresses fragmentation and sharing across requests; it likewise does not change a single request’s compute.

Getting this order backwards is common: reach for quantisation and paging first, observe that time-to-first-token has not moved at all, and conclude the hardware is inadequate. Establish which term binds before choosing a tool. The measurement is simple — record time-to-first-token and the mean inter-token interval separately, and see which dominates total time.

For estimating VRAM itself — how weights, KV cache and fixed overhead decompose, and the constraint that KV head count must divide evenly across GPUs in tensor-parallel setups — see Turning Three Known VRAM Traps Into a Working Estimator, which gives a directly applicable formula along with its domain of validity.

Production Pitfalls & Engineering Traps

1. FlashAttention Limitations with Dense Visual Sequences

While FlashAttention-2 is a staple for reducing VRAM in standard LLMs, in MLLMs, visual tokens are often densely packed at the very beginning of the prompt sequence. If the sheer volume of visual tokens from cross-frames or gigapixel images exceeds FlashAttention’s block size, VRAM fragmentation and OOM can still occur. Mitigation: Implement PagedAttention (via frameworks like vLLM) for contiguous visual tokens to manage VRAM at the page block level, completely avoiding memory fragmentation.

2. Visual Token Context Length Management

In multi-turn multimodal chats, developers frequently make the mistake of concatenating raw visual tokens from historical images into the active prompt. This practice pushes the KV Cache rapidly towards the model’s maximum context window limit (e.g., 4K or 8K). Mitigation: Do not cache raw historical visual features across multiple turns. Summarize previous visual inferences into text caching, or introduce Token Merging (ToMe) post-projector to systematically drop redundant background visual tokens.

Leave a Reply

Scroll down