[KongchangAI]
· 2 min read· 1,215 words

One Parameter Fits 128K Context on a 4090: FP8 KV Cache Quantization in Practice

One Parameter Fits 128K Context on a 4090: FP8 KV Cache Quantization in Practice

vLLM's FP8 KV cache quantization halves VRAM, letting 64K context run on a 24GB GPU, but 128K still needs other optimizations.

The article points out that the real bottleneck for long-context inference is the KV cache, which grows linearly with context, not the fixed-size model weights. Taking Llama 3.1 8B as an example, FP16 weights take about 16GB; at 32K context the 4GB KV cache is fine, but at 64K the 8GB cache pushes total VRAM to 26GB—over the 24GB limit—and at 128K single-GPU deployment is impossible. vLLM offers the `--kv-cache-dtype fp8` parameter, compressing the KV cache from 2 bytes to 1 byte. At 64K the cache drops to 4GB for a total of about 22GB, just fitting a 4090 and cutting cost from $1.5/hour to $0.35/hour. FP8 should use E4M3 rather than E5M2, since the latter's lower mantissa precision degrades attention scores. However, 128K context with FP8 cache alone still exceeds 24GB, requiring FP8 weights or a more GQA-friendly model like Qwen2.5 7B. Production deployment requires attention to hardware support, accuracy calibration, and retrieval-task evaluation.

Why KV Cache Is the Real Bottleneck for Long Contexts

Many people think that getting a 7B model to run on a 24GB GPU is all there is to it—until they try a 128K context window and watch VRAM explode. According to calculations from the developer community, Llama 3.1 8B takes about 16GB with FP16 weights, plus 2GB of runtime overhead. At 32K context, the KV cache takes 4GB for a total of about 22GB, which is still comfortable. But at 64K, the cache doubles to 8GB and the total hits 26GB—already over the 24GB red line of the 4090. At 128K, the cache balloons to 16GB and the total is about 34GB, making single-GPU deployment completely out of reach.

The essence of the problem: weights are fixed, but the KV cache grows linearly with context length. So instead of messing with weight precision, it's better to attack the cache directly—which is also the lowest-cost step.

rss source: Your 128K Context Fits on a 4090 — If You Quantize the KV Cache

KV cache is the mechanism that stores the Key and Value tensors of historical tokens during Transformer autoregressive inference. When generating each new token, the model needs to attend to the Keys and Values of all previous tokens. Caching them avoids redundant computation, reducing inference complexity from O(n²) to O(n). But cache size is proportional to context length, number of layers, number of attention heads, and head dimension: for Llama 3.1 8B (32 layers, 32 heads, head dimension 128), the KV cache per token is about 0.5MB (FP16), which at 32K context would be about 16GB? Actual calculation: 2 (K and V) × 32 layers × 32 heads × 128 dims × 2 bytes = 512KB/token, so 32K is about 16GB, which doesn't match the 4GB in the article. This suggests the article may be calculating with GQA (Grouped Query Attention). Llama 3.1 8B uses GQA with 8 KV heads, so the per-token cache is 2×32×8×128×2=128KB, and 32K is about 4GB—which matches.

One Command-Line Parameter, Cache Halved

In vLLM, you just add --kv-cache-dtype fp8. Weights stay at FP16/BF16, and the KV cache is compressed from 2 bytes per value to 1 byte. According to the official vLLM documentation, fp8 corresponds to the E4M3 format and offers three calibration modes: no calibration (scaling factor defaults to 1.0), random-token warmup, and dataset calibration via llm-compressor (recommended when pursuing maximum accuracy).

The effect is immediate: at 64K context, the FP16 cache drops from 8GB to 4GB, for a total of about 22GB—just enough to fit on a 4090. A 7B/64K workload that previously required an 80GB GPU at $1.5+/hour can now run on a single 24GB 4090 at about $0.35/hour, cutting cost to roughly a quarter—at the price of just one CLI parameter.

Why FP8 Is Particularly Suited to KV Cache

Weight quantization steals most of the attention, but KV values are actually an easier target: within a single attention head, values are concentrated in a narrow, well-behaved range. E4M3's narrow dynamic range plus 3 mantissa bits matches this distribution perfectly. E5M2, despite its wider dynamic range, wastes most of that range on KV cache, and its 2 mantissa bits noticeably degrade attention scores—an error that compounds and amplifies across layers.

In other words, format choice itself is a technical decision: the default should be fp8 (i.e., E4M3), not fp8_e5m2. This is an easily overlooked detail with far-reaching impact.

The Math Doesn't Lie: 128K Still Doesn't Fit

To be honest: with FP8 cache alone, Llama 3.1 8B still can't run 128K on a 24GB card. 8GB cache plus 16GB weights plus 2GB overhead totals about 26GB—still over.

There are two ways out: first, stack FP8 weights on top, giving 8+8+2 ≈ 18GB, which fits easily; second, switch to an architecture with more aggressive GQA, such as Qwen2.5 7B, whose FP16 cache at 128K is only about 7GB, or about 3.5GB with FP8, for a total of about 19GB. This parameter is powerful, but it's not magic—the arithmetic rules still reign.

ContextFP16 Cache Total VRAMFP8 Cache Total VRAM
32K4GB → about 22GB ✅2GB → about 20GB ✅
64K8GB → about 26GB ❌4GB → about 22GB ✅ fits 4090
128K16GB → about 34GB ❌8GB → about 26GB ❌ still over

Three Things to Watch Before Going Live

Hardware support differences: The biggest gains appear on Hopper-architecture GPUs with native FP8 support. Memory savings on older cards are real, but you need to confirm kernel support in your vLLM version—the compute path may not be equally optimized.

Don't skip accuracy calibration: The default scaling factor of 1.0 is fine to start with, but for production, dataset calibration via llm-compressor is recommended. If you want finer-grained per-head quantization than per-tensor, you'll also need the FlashAttention backend to work with the calibration process.

Evaluate retrieval capability, not perplexity: Accuracy drops on standard benchmarks are usually small, but exact-recall long-context tasks (such as needle-in-a-haystack document QA) will expose problems first. Be sure to run your own evaluation set before and after deployment.

The Trade-off Logic in Practice

The real value of this approach is keeping people on cheaper cards. Quantizing the cache isn't an isolated trick—it combines with model architecture choice and weight precision: up to 64K, FP8 cache alone is enough; for 128K, you need FP8 cache plus FP8 weights, or simply switch to a model with a more GQA-friendly architecture.

The author also built a free GPU Advisor tool: enter model size, context length, and quantization tolerance, and it recommends the cheapest viable configuration. But for most teams, validating FP8 KV cache accuracy on your own workload is probably more convincing than any recommendation tool.

Background

FP8 E4M3 is an 8-bit floating-point format with 1 sign bit, 4 exponent bits, and 3 mantissa bits. Its dynamic range is about ±448, with roughly 2 decimal significant digits of precision. Compared with FP16 (5 exponent bits, 10 mantissa bits), E4M3 sacrifices dynamic range and precision but is better suited to scenarios with concentrated value distributions. In the KV cache, Key and Value values within each attention head typically fluctuate within a narrow range, which E4M3's narrow dynamic range covers exactly, while its 3 mantissa bits retain sufficient precision. E5M2 has 5 exponent bits and 2 mantissa bits, offering a wider dynamic range (capable of representing very large and very small values), but its lower mantissa precision leads to larger errors in attention score computation—errors that compound across layers. That's why vLLM recommends fp8 (E4M3) by default rather than fp8_e5m2.

llm-compressor is a model quantization library in the vLLM ecosystem. It supports FP8, INT8, and other quantization for weights and activations, and provides dataset-based calibration workflows. Calibration works by feeding representative data (e.g., domain text), measuring tensor dynamic ranges, and computing optimal scaling factors to minimize quantization error. For KV cache, calibration can be done separately for Keys and Values, and even refined to per-head granularity—but this requires the FlashAttention backend to support the corresponding quantization kernels. No calibration (scaling factor 1.0) works in most cases, but if the data distribution shifts significantly, it may lead to overflow or accuracy loss.

#FP8 KV cache#vLLM#KV cache quantization#long context inference#24GB GPU#E4M3#llm-compressor
Share:

Related articles