LLM Inference Optimization Practice: The Path from OOM Crash to Stable 3GB Memory Operation

Original Author/Maintainer: Alcimarrfilho, Source Platform: GitHub, Original Link: https://github.com/Alcimarrfilho/llm-inference-optimization

A detailed LLM inference optimization experiment report showing how to optimize 16K context inference from a 31GB VRAM OOM error to stable 3GB operation using QLoRA, KV Cache, and SDPA technologies, and discussing State Space Models (e.g., Mamba) as a future direction for ultra-long context expansion.

In practical LLM applications, resource consumption during inference is a key challenge, and processing long contexts (e.g., 16K tokens) faces severe VRAM bottlenecks. This experiment records the complete process from OOM crash to successful optimization, providing practical experience for developers.

Experimental Environment: Google Colab T4 GPU (15GB VRAM), Test Model: TinyLlama-1.1B-Chat-v1.0, Objective: Achieve stable 16K context inference under limited hardware.

The experiment uses three complementary technologies to solve VRAM issues:

1. QLoRA: 4-bit quantization compresses model weights, reducing static VRAM usage while maintaining performance with low-rank adapters;
2. KV Cache: Caches previously computed Key and Value vectors to avoid redundant calculations, reducing the time complexity of self-attention from O(n²) to O(n);
3. SDPA: PyTorch's fused attention implementation, which avoids materializing the full attention matrix through block-wise computation, suitable for T4 GPUs (which do not support FlashAttention-2).

Performance data for three key stages:

• Stage 1: Loading the model with QLoRA 4-bit quantization, VRAM usage: 805.93 MB;
• Stage 2: Processing 16K tokens without optimization, requiring 30.91 GB VRAM leading to OOM;
• Stage 3: Combined KV Cache + SDPA optimization, generation time: 4.13 seconds, peak VRAM: 3055.28 MB (≈3GB), achieving over 90% VRAM savings.

Although the Transformer architecture has been optimized, KV Cache VRAM still grows linearly with sequence length (e.g., 2 million tokens require hundreds of GB). State Space Models (SSM) like the Mamba architecture provide a solution:

• Does not store the full token history; compresses information into a fixed-size hidden state with O(n) memory complexity;
• Core advantages: Selective state space, hardware-aware algorithms, linear inference time, suitable for ultra-long context scenarios.

Key insights from this experiment:

1. Quantization is the first step in VRAM optimization: 4-bit quantization compresses the model size to 1/4;
2. KV Cache is essential for long sequence generation: Reduces redundant calculations and lowers latency;
3. Attention optimization needs to adapt to hardware: FlashAttention-2 is optimal but not supported by all GPUs; SDPA has good compatibility;
4. Architecture choice determines expansion limits: For ultra-long context needs, consider SSM architectures like Mamba.

Complete reproduction path:

1. Environment Preparation: Open laboratorio_10.ipynb and run it in Google Colab;
2. Hardware Configuration: Python3 environment + T4 GPU;
3. Dependency Installation: The notebook automatically installs transformers, bitsandbytes, and accelerate libraries;
4. Sequential Execution: Execute the code cells in order to reproduce the experiment.