Skip to content

JohnScheuer/long-context-benchmark

Repository files navigation

long-context-benchmark

Author: João Felipe De Souza

Python PyTorch Transformers bitsandbytes CUDA Platform GPU License


Overview

Long-context inference benchmark pushing Qwen2-0.5B from 4K to 32K tokens on an RTX 2070 with 8.6 GB VRAM.

Three optimization layers tested:

  1. SDPA attention (memory-efficient kernel)
  2. Chunked prefill (reduces peak activation memory)
  3. INT4 quantization (reduces model weight memory)

For architecture details, see DESIGN.md.


Why This Matters

Long-context inference hits a capacity cliff on consumer GPUs. Without optimizations, context beyond 4K causes VRAM overflow, VRAM swap, and catastrophic performance collapse.

This benchmark maps exactly:

  • where the cliff is
  • which optimizations move it
  • which ones don't help (and why)

Model and Hardware

Model

  • Qwen2-0.5B (494M parameters, RoPE, supports 32K context)

Hardware

  • NVIDIA RTX 2070 (8.6 GB VRAM)
  • CUDA 13.0

Configurations tested

  • Context lengths: 512, 1024, 2048, 4096, 8192, 16384, 32768
  • Batch sizes: 1, 2, 4, 8
  • Precision: FP16, INT4 (NF4)
  • Attention: eager (vanilla), SDPA (memory-efficient)
  • Prefill: full, chunked (1024-token chunks)

Key Findings

Finding 1 — SDPA + chunked prefill gives 40x speedup at 8K

Metric Vanilla (eager) SDPA + chunked
ctx=8192 prefill 40854 ms 1022 ms
ctx=8192 peak VRAM 9553 MB 1642 MB
Speedup 40×
Memory reduction 83%

The vanilla result was VRAM swap (9.5 GB on an 8.6 GB GPU). SDPA + chunked prefill eliminated the swap entirely.

Finding 2 — 32K context now works on RTX 2070

Context FP16 Prefill FP16 Peak FP16 Throughput
2048 1485 ms 5905 MB 205.0 tok/s (bs=8)
4096 1429 ms 3524 MB 99.0 tok/s (bs=4)
8192 1667 ms 3716 MB 48.0 tok/s (bs=4)
16384 4698 ms 2525 MB 25.3 tok/s (bs=2)
32768 7536 ms 2001 MB 12.6 tok/s (bs=1)

Previously, 8K+ was impossible on this GPU.

Finding 3 — FP16 beats INT4 at long context

Context FP16 INT4
8192 bs=1 8139ms, 1642MB 3366ms, 1967MB
16384 bs=1 2371ms, 1738MB 11666ms, 3180MB
32768 bs=1 7536ms, 2001MB OOM

INT4 uses less model weight memory (430MB vs 942MB) but:

  • KV-cache is still FP16 in both cases
  • INT4 adds dequantization overhead that grows with context
  • At long context, KV-cache dominates, not model weights

Finding 4 — The capacity cliff moves dramatically with optimization

Max batch per context:

Context Vanilla SDPA + chunked FP16 INT4
2048 8 8 8
4096 2 4 4
8192 OOM 4 2
16384 OOM 2 1
32768 OOM 1 OOM

Finding 5 — Prefill scaling is no longer quadratic

With chunked prefill, prefill latency scales more linearly:

Context Vanilla Prefill SDPA + chunked Prefill
2048 bs=1 524 ms 990 ms
4096 bs=1 1681 ms 459 ms
8192 bs=1 40854 ms (swap) 1022 ms
16384 bs=1 OOM 2371 ms
32768 bs=1 OOM 7536 ms

Main Conclusion

Long-context inference on consumer GPUs requires a stack of optimizations:

  1. SDPA efficient attention — reduces activation memory from O(n^2) to O(n)
  2. Chunked prefill — limits peak memory per forward pass
  3. GQA (from multi-query-attention project) — reduces KV-cache by 2-12x

Quantization (INT4) is NOT a long-context solution. It reduces model weight memory but doesn't help with the real bottleneck: KV-cache and activations.

The practical long-context stack:

SDPA efficient → chunked prefill → GQA → KV-cache eviction

NOT:

quantize everything → hope it fits

Results Files

File Description
results/long_context_final.csv Baseline (vanilla + eager)
results/long_context_better.csv SDPA + chunked prefill
results/long_context_ultra.csv FP16 vs INT4 comparison
results/metadata.json Configuration

Plots

File Description
plots/prefill_vs_context.png Prefill latency scaling
plots/throughput_vs_context.png Decode throughput
plots/memory_vs_context.png VRAM usage
plots/max_batch_capacity.png Max batch per context
plots/capacity_heatmap.png OK/OOM capacity map

Repository Structure

long-context-benchmark/
├── long_context_benchmark.py      # Initial vanilla benchmark
├── long_context_safe.py           # Memory-safe version
├── long_context_final.py          # Calibrated estimator version
├── long_context_better.py         # SDPA + chunked prefill
├── long_context_ultra.py          # FP16 vs INT4 comparison
├── plot_long_context.py           # Visualization
├── README.md
├── DESIGN.md
├── LICENSE
├── requirements.txt
├── results/
└── plots/

How to Run

1. Setup

python3 -m venv venv && source venv/bin/activate
pip install -r requirements.txt

2. Baseline benchmark

python3 long_context_final.py

3. SDPA + chunked prefill benchmark

python3 long_context_better.py

4. FP16 vs INT4 ultra benchmark

python3 long_context_ultra.py

5. Generate plots

python3 plot_long_context.py

How This Connects to Other Projects

Problem Solution project
Attention O(n^2) memory flash-attention-benchmark
KV-cache too large kv-cache-eviction-benchmark
Prefill too slow prefill-chunking-profiler
Model weights too large quantization-profiler
Decode too slow decode-batching-profiler
Memory budget unknown gpu-memory-profiler
PE affects eviction rope-vs-absolute-pe-benchmark

Limitations

  • Single model (Qwen2-0.5B)
  • Single GPU (RTX 2070)
  • INT4 KV-cache not tested (would further extend context)
  • No KV-cache eviction applied
  • SDPA efficient may not use fused kernels for GQA on Turing (mismatched head counts)

References

  • Kwon et al., PagedAttention (2023)
  • Dao et al., FlashAttention (2022)
  • Xiao et al., Attention Sinks (2023)
  • Dettmers et al., QLoRA (2023)

About

Long-context benchmark pushing Qwen2-0.5B from 4K to 32K tokens on RTX 2070 using SDPA + chunked prefill. Shows 40x speedup at 8K, FP16 beating INT4 at long context, and that quantization is NOT a long-context solution — KV-cache is the real bottleneck.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages