Log memory usage

. The data proves that intelligent memory management and quantization selection are multiplicative factors in inference efficiency.

Core Solution

Effective GPU memory management requires a layered approach: accurate profiling, quantization selection, and architectural optimizations like PagedAttention.

1. Memory Profiling and Baseline

Before optimization, establish a memory baseline. Use torch.cuda.memory_summary() or vendor tools to identify weight vs. activation vs. cache usage.

import torch
import transformers

model_id = "meta-llama/Llama-2-7b-hf"
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Log memory usage
print(f"Peak Memory: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")
torch.cuda.reset_peak_memory_stats()

2. Quantization Implementation

Select quantization based on SLA requirements. AWQ is recommended for most production inference due to its accuracy-speed balance. Use auto-gptq or llm-awq for weight conversion.

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load AWQ quantized model
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-2-7B-Chat-AWQ",
    torch_dtype=torch.float16,
    device_map="auto"
)

3. PagedAttention Architecture

For high-throughput serving, deploy vLLM. vLLM implements PagedAttention, which manages KV cache as virtual memory pages. This eliminates fragmentation and enables dynamic batching, maximizing GPU utilization.

Architecture Decision:

• Rationale: Traditional implementations allocate contiguous memory for KV cache, leading to internal fragmentation. PagedAttention divides KV cache into blocks, allowing non-contiguous allocation similar to OS paging. This reduces waste and allows the system to pack more requests into the same VRAM.
• Implementation: Configure gpu_memory_utilization to reserve memory for KV cache blocks while leaving headroom for activations.

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    quantization="awq",
    gpu_memory_utilization=0.90,  # Reserve 10% for overhead
    max_model_len=4096,           # Cap context to control KV cache growth
    block_size=16,                # PagedAttention block size
    enforce_eager=False           # Use CUDA graphs for latency reduction
)

outputs = llm.generate("Explain GPU memory management.", SamplingParams(temperature=0.7))

4. KV Cache Optimization

KV cache memory scales with batch_size * sequence_length * num_layers * hidden_size.

• Strategy: Limit max_model_len to the practical maximum required by your application.
• Strategy: Enable use_cache=True but monitor cache eviction policies if serving streaming requests with variable lengths.
• Strategy: In vLLM, tune block_size. Smaller blocks reduce waste for short sequences but increase metadata overhead. 16 is optimal for most workloads.

Pitfall Guide

1. Ignoring KV Cache Growth

Mistake: Allocating memory for weights but failing to account for KV cache expansion during long generations. Impact: OOM errors occur mid-generation when context exceeds reserved space. Best Practice: Calculate KV cache footprint: 2 bytes * batch * seq_len * layers * hidden. Reserve 30-50% of VRAM for cache in dynamic workloads.

2. Over-Quantization Without Validation

Mistake: Applying INT4 quantization to all models regardless of domain. Impact: Significant accuracy degradation in specialized domains (code, math, medical) where outlier weights carry critical information. Best Practice: Benchmark perplexity on domain-specific datasets post-quantization. Use AWQ to preserve outliers. Revert to INT8 or FP16 if degradation exceeds SLA thresholds.

3. Fragmentation in Naive Batching

Mistake: Using static batch sizes with fixed memory allocation per request. Impact: VRAM fragmentation leaves unusable gaps, reducing effective capacity by 20-40%. Best Practice: Use PagedAttention-based engines (vLLM, TGI). Avoid custom serving loops that allocate contiguous tensors per request.

4. CPU Offloading Bottlenecks

Mistake: Relying on CPU offloading to fit models that exceed GPU VRAM. Impact: PCIe bandwidth limits throughput to <10 tok/s, making the system unusable for interactive applications. Best Practice: CPU offloading should only be used for cold-start or non-latency-sensitive batch jobs. For interactive use, reduce model size or quantization level to fit entirely in VRAM.

5. Static GPU Memory Utilization Caps

Mistake: Hardcoding gpu_memory_utilization=0.8 without profiling. Impact: Either underutilization (leaving 20% VRAM idle) or instability (OOM due to activation spikes). Best Practice: Profile peak memory usage under load. Set utilization to 0.90 for stable workloads with predictable batch sizes. Use dynamic adjustment in auto-scaling groups.

6. CUDA Graph Incompatibility

Mistake: Enabling CUDA graphs with variable input shapes or unsupported operators. Impact: Graph capture fails or falls back to eager mode, increasing latency. Best Practice: Ensure fixed input shapes where possible. Verify operator support in the quantization backend. Use vLLM's automatic CUDA graph capture which handles shape padding safely.

7. Neglecting Activation Memory

Mistake: Focusing solely on weight memory and ignoring activation memory during forward passes. Impact: OOM during peak compute phases, especially with large batch sizes. Best Practice: Monitor torch.cuda.max_memory_reserved(). Activation memory scales with batch size; reduce batch size if activation spikes cause OOMs.

Production Bundle

Action Checklist

• Profile baseline VRAM usage for weights, KV cache, and activations.
• Select quantization method (AWQ recommended) and validate perplexity on domain data.
• Deploy vLLM with PagedAttention to eliminate fragmentation.
• Configure gpu_memory_utilization to 0.90 after profiling peak usage.
• Set max_model_len to the maximum practical context length to cap KV cache.
• Enable CUDA graphs for latency reduction if input shapes are stable.
• Implement monitoring for KV cache usage and OOM error rates.
• Test worst-case scenario: max batch size + max context length.

Decision Matrix

Scenario	Recommended Approach	Why	Cost Impact
Low latency, high accuracy SLA	FP16 + Tensor Parallel	Maximizes compute efficiency; avoids quantization overhead.	High GPU count required.
High throughput, cost-sensitive	INT4 AWQ + vLLM	Maximizes density; PagedAttention enables dynamic batching.	Low GPU count; optimal TCO.
Single consumer GPU (24GB)	INT4 GGUF + llama.cpp	GGUF offers CPU/GPU hybrid offloading; llama.cpp optimized for consumer hardware.	Zero infrastructure cost.
Long context (>32K tokens)	FP16 + FlashAttention-2 + KV Cache Pruning	FlashAttention reduces memory bandwidth pressure; pruning manages cache growth.	Moderate GPU cost; requires complex setup.
Multi-tenant serving	vLLM + Continuous Batching	Isolates requests; maximizes utilization via dynamic batching.	Scales efficiently; reduces per-request cost.

Configuration Template

Copy this configuration for a production vLLM deployment with AWQ quantization:

# vllm_config.yaml
model: "TheBloke/Llama-2-7B-Chat-AWQ"
quantization: "awq"
gpu_memory_utilization: 0.90
max_model_len: 4096
block_size: 16
enforce_eager: false
max_num_batched_tokens: 4096
max_num_seqs: 256
trust_remote_code: true
dtype: "float16"
swap_space: 4  # CPU swap space in GB for overflow protection

Docker Run Command:

docker run --gpus all -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  vllm/vllm-openai:latest \
  --model TheBloke/Llama-2-7B-Chat-AWQ \
  --quantization awq \
  --gpu-memory-utilization 0.90 \
  --max-model-len 4096 \
  --block-size 16

Quick Start Guide

1. Install vLLM:

   pip install vllm autoawq
   

2. Launch Server:

   python -m vllm.entrypoints.openai.api_server \
     --model meta-llama/Llama-2-7b-chat-hf \
     --quantization awq \
     --gpu-memory-utilization 0.90
   

3. Verify Memory: Check logs for GPU KV cache usage. Ensure it is within limits. Use nvidia-smi to confirm VRAM usage matches expectations (~5-6 GB for 7B AWQ).
4. Benchmark: Run a load test with wrk or locust targeting the /v1/completions endpoint. Monitor throughput and latency. Adjust max_num_seqs based on observed utilization.
5. Optimize: If latency is high, enable CUDA graphs by ensuring enforce_eager is false and input shapes are consistent. If OOM occurs, reduce max_model_len or gpu_memory_utilization.

GPU memory management for LLMs is not a one-time configuration but a continuous optimization process. By leveraging PagedAttention, selecting appropriate quantization, and rigorously profiling memory components, engineers can achieve significant gains in throughput and cost efficiency while maintaining model quality.