Sari la conținut

Română · beta. Traducerea a fost făcută automat și este încă în curs de revizuire. Citește pagina în engleză

perspectiva aiqbenchCipuri & infrastructură

analiza aiqbench

This article details a technical integration of 4-bit KV quantization with hierarchical caching on AMD MI355X hardware to manage memory constraints for long-context AI agents. By offloading quantized key-value states from HBM to CPU DRAM, the approach aims to increase throughput and reduce latency without sacrificing model accuracy.
Sursă care acoperăAMD
Ce este aceastaLectura unui model asupra articolului și a acoperirii conexe, scrisă pentru aiqbench. Nu sunt cuvintele publicației.

4-bit KV Caching in LMCache: Offloading Quantized KV Beyond HBM for Context-Heavy Agents on AMD MI355X

AMD7 min de cititArticolul și perspectiva sunt în engleză.

As agents carry ever-longer context from turn to turn, the KV cache becomes the resource that runs out first. Two techniques ease that pressure from different angles: KV quantization makes each cached token smaller, while a hierarchical cache like LMCache spills cold KV to CPU DRAM. The two have mostly been developed independently.

13 min read. | 3138 total words.

As agents carry ever-longer context from turn to turn, the KV cache becomes the resource that runs out first. Two techniques ease that pressure from different angles: KV quantization makes each cached token smaller, while a hierarchical cache like LMCache spills cold KV to CPU DRAM. The two have mostly been developed independently.

This blog brings them together on AMD Instinct™ MI355X: a KV cache that is both 4-bit and tiered, taking quantized KV (TurboQuant) past the HBM boundary into CPU DRAM and back, bit-exact, accuracy-neutral, and fast. We walk through the layout-aware connector that makes a mixed bf16 + 4-bit cache safe to offload, the serving methodology we used to measure it, and what happens to goodput, latency, and accuracy when quantization and tiering compound at equal memory.

LMCache extends TurboQuant past the HBM wall. 4-bit KV holds more, but big agent contexts still overflow HBM. LMCache catches the overflow in CPU DRAM and reloads it fast, so a miss is a quick copy, not a full re-prefill. That turns 4-bit's saved space into real goodput (more requests served), not just a bigger cache.

LMCache extends TurboQuant past the HBM wall. 4-bit KV holds more, but big agent contexts still overflow HBM. LMCache catches the overflow in CPU DRAM and reloads it fast, so a miss is a quick copy, not a full re-prefill. That turns 4-bit's saved space into real goodput (more requests served), not just a bigger cache.

Quantized KV offload with TurboQuant. LMCache now stores, evicts, and reloads a mixed bf16 + 4-bit KV stack through CPU DRAM, using our optimized TurboQuant 4-bit codebook, the AMD-productionized kernels from turboquant-vllm-agentic.

Quantized KV offload with TurboQuant. LMCache now stores, evicts, and reloads a mixed bf16 + 4-bit KV stack through CPU DRAM, using our optimized TurboQuant 4-bit codebook, the AMD-productionized kernels from turboquant-vllm-agentic.

A layout-aware connector. TurboQuant doesn't store KV in one uniform format: it keeps the sensitive boundary layers in bf16 and packs the middle layers into 4-bit. The LMCache connector understands the mix and moves every byte back exactly.

A layout-aware connector. TurboQuant doesn't store KV in one uniform format: it keeps the sensitive boundary layers in bf16 and packs the middle layers into 4-bit. The LMCache connector understands the mix and moves every byte back exactly.

Quantization makes the cache fit. At a fixed HBM+DRAM budget, TurboQuant keeps ~3.2× more KV resident in HBM and lifts the LMCache hit rate from 75.2% (BF16) to 86.8% (TurboQuant), so most reused context is served straight from cache instead of being recomputed.

Quantization makes the cache fit. At a fixed HBM+DRAM budget, TurboQuant keeps ~3.2× more KV resident in HBM and lifts the LMCache hit rate from 75.2% (BF16) to 86.8% (TurboQuant), so most reused context is served straight from cache instead of being recomputed.

Big serving win. TurboQuant delivers ~2.0× the goodput — requests completed in the fixed 20-min replay — at ~2.6× lower p95 TTFT than a BF16 cache, while moving ~1.7× less data across at equal memory.

Big serving win. TurboQuant delivers ~2.0× the goodput — requests completed in the fixed 20-min replay — at ~2.6× lower p95 TTFT than a BF16 cache, while moving ~1.7× less data across at equal memory.

An agent carries the same context from turn to turn: system prompt, tools, retrieved files, a running plan. That's tens to hundreds of thousands of tokens, ~95% of it reused each turn. At those lengths it's the KV cache, not the model weights, that fills HBM. And once HBM is full, the cache has to evict old prefixes, so every reused turn hits one of two bad options: re-prefill the dropped prefix (pay its full compute and TTFT again) or reload it from a slower tier.

LMCache is an open-source KV-cache layer for vLLM that stores and reuses KV across requests, so repeated context is reloaded instead of recomputed. It spreads that KV across a memory hierarchy ordered by speed and size: GPU HBM (fast, small) → CPU DRAM (slower, large) → disk (slowest, largest). Each level is a tier, and LMCache reloads KV from a lower tier on demand instead of dropping it (see the LMCache blog). But by default, LMCache keeps KV at full precision (bf16), which is heavy: HBM fits only a few tokens, fills up and evicts early, and every reload from DRAM copies a lot of bytes, slowly. The result is a cache that spills to DRAM early and keeps moving big blocks back.

Quantization attacks both problems at once. Packing the KV cache to 4-bit makes each token ~3.76× smaller, which buys two things: ~3.2× more tokens stay in fast HBM before anything spills to DRAM, and tokens that do spill move as 4-bit datatypes, a much smaller payload. These two effects compound: more reuse is served from HBM, and the spills that remain are cheaper.

vLLM serves the model through its KV-connector hook, handing each request's KV to LMCacheConnectorV1, which tiers GPU HBM and CPU DRAM, evicting and reloading on demand. What flows through that path is a mixed bf16 + 4-bit KV stack instead of uniform bf16.

We don't compress every layer, and that's deliberate. The first two and last two attention layers are the most accuracy-sensitive, so they stay bf16; the middle 58 layers tolerate compression, so they're packed to 4-bit. This keeps full precision where it matters and takes memory savings everywhere else (vLLM TurboQuant OSS). One model holds two physically different KV tensors:

Our implementation is on top of LMCache (GPU connector, vLLM integration, kernels, correctness tests) and builds on LMCache's existing layer-group support to make the connector group-aware for a quantized workflow. At registration it partitions layers into two transfer groups (bf16, packed-uint8) and issues one correctly-strided copy per group, so both ride the same HBM-to-DRAM path. Two properties keep it a no-config run:

Auto-selection. A uint8 KV container is an unambiguous packed signal, so the group-aware path turns on by itself — no new flags, and one code path regardless of how the 4-bit slot is packed.

Auto-selection. A uint8 KV container is an unambiguous packed signal, so the group-aware path turns on by itself — no new flags, and one code path regardless of how the 4-bit slot is packed.

Byte-transparency. The slot's internal K/V split never leaks into the cache; LMCache copies whole slots as opaque uint8, so the round-trip is byte-exact however K and V are packed inside.

Byte-transparency. The slot's internal K/V split never leaks into the cache; LMCache copies whole slots as opaque uint8, so the round-trip is byte-exact however K and V are packed inside.

We run on 2× AMD Instinct MI355X GPUs (gfx950, 288 GB HBM each), which gives us a large but finite HBM budget to put under pressure. The serving stack is ROCm 7.2.3 with vLLM V1, and LMCache built for HIP and extended with our group-aware connector. Our test model is MiniMax-M2.5, a 230 GB FP8 mixture-of-experts, served with tensor parallelism across both GPUs (TP=2). As a safety measure, we also set PYTHONHASHSEED=0 and LMCACHE_PRE_CACHING_HASH_ALGORITHM=sha256 on every run to keep cache keys stable across TP workers.

KV encodings. BF16 (baseline) vs TurboQuant: our AMD-productionized 4-bit build with a custom FlyDSL kernel.

Three cache configurations. We run the same workload three times, swapping only the KV strategy (see LMcacheMI300X for reference and settings):

Every benchmark replays the same agentic trace for 20 minutes: many users hit the server at once, and each finished conversation is replaced by a new one, keeping the server continuously loaded. Terms:

Agentic trace replay: we re-issue 739 recorded, anonymized Claude Code sessions (multi-turn, tool-using coding), so the load is real agent traffic, not synthetic.

Text integral sub licența deschisă a publicației · 1285 cuvinte

Subiecte conexe