abhinandan

Artifact 16

Building a Quantized L2 KV Cache for SGLang - HiQCache

October 12, 20264views

one of the most useful optimizations in LLM inference is KV caching. Instead of recomputing the same prefix every time, we store its previously calculated keys and values and reuse them when needed.

But there's a problem. As context lengths grow and more requests arrive, the KV cache also grows. GPU memory is limited, so we cannot keep every reusable prefix in VRAM forever.

This is where hierarchical KV caching comes in. Rather than completely discarding KV that no longer fits on the GPU, we can keep a copy in CPU memory and bring it back whenever a future request needs it.

SGLang already does this through HiCache. But while exploring its implementation, one question interested me: if the CPU copy is primarily used for storage, do we really need to keep it in the same precision as the GPU copy?

For example, suppose attention needs BF16 KV on the GPU. We could keep the active KV in BF16, but store the CPU copy in a smaller INT8 representation. Whenever that KV is needed again, we could restore it to BF16 before attention uses it.

That became the idea behind HiQCache, an experimental quantized L2 KV-cache implementation for SGLang.

I wanted to explore whether I could make the existing CPU cache more memory-efficient without changing the attention kernels, caching policy, or scheduler.

The result was 43.75% less host KV memory per token, giving 1.78× the L2 token capacity at the same CPU memory budget. In a shared-prefix workload, this also increased cache reuse.

But getting there involved more than quantizing a tensor. The interesting challenges were fitting a compressed representation into SGLang's existing memory layout, using its CUDA transfer kernels, and making sure asynchronous restoration remained correct.

First, how does hierarchical KV caching work?

Before getting into compression, let's understand what is actually being cached.

Suppose multiple requests share a long system prompt. When we process that prefix for the first time, the model computes the corresponding keys and values across its transformer layers.

When another request starts with the same prefix, we can reuse that KV instead of running the same prefix through the model again.

That's the basic idea of prefix caching.

Now suppose we have thousands of reusable prefixes. Keeping all of them in GPU memory can become expensive, especially when we're also using that memory for model weights, activations, and active requests.

So instead of relying on GPU memory alone, SGLang's HiCache introduces multiple storage tiers:

text
                 SGLang HiCache

          ┌──────────────────────┐
          │ L1 — GPU memory      │
          │ Active KV cache      │
          └──────────┬───────────┘
                     │
                Backup / Restore
                     │
          ┌──────────▼───────────┐
          │ L2 — CPU memory      │
          │ Offloaded KV cache   │
          └──────────┬───────────┘
                     │
          ┌──────────▼───────────┐
          │ L3 — External        │
          │ Optional storage     │
          └──────────────────────┘

A small clarification: L1 and L2 here are logical KV-cache tiers, not the physical L1 and L2 caches inside the GPU. HiCache's L1 lives in GPU VRAM, while its L2 lives in CPU DRAM.

When reusable KV is backed up to L2, HiCache can retain it even after the corresponding GPU cache entries are evicted. On a later prefix-cache hit, that KV can be restored to the GPU instead of recomputing the prefix.

Of course, restoring KV from CPU memory has a transfer cost. But for sufficiently expensive prefixes, transferring already-computed KV can be cheaper than recomputing it.

There is still one limitation: the CPU cache also has a finite capacity.

And that is the part I wanted to optimize.

Can we make the CPU KV cache smaller?

In the normal BF16 HiCache configuration, the CPU representation mirrors the GPU KV precision.

Let's take Qwen3-8B, the model I used for the experiments.

It has:

  • 36 transformer layers
  • 8 KV heads
  • 128 dimensions per KV head
  • BF16 KV values, which occupy 2 bytes each

For one token, across all layers, the total KV storage is:

text
KV bytes per token
    = 2 × layers × KV heads × head dimension × bytes/value
    = 2 × 36 × 8 × 128 × 2
    = 147,456 bytes

The first 2 represents keys and values.

So one cached token consumes 144 KiB across the model's layers. A long prefix containing thousands of tokens can occupy hundreds of megabytes.

With an 8 GB CPU cache budget, the tested BF16 configuration accommodates approximately 54,254 tokens.

Now what if we store the CPU representation in INT8 instead?

INT8 uses one byte per value instead of two. We cannot simply cast BF16 values to integers, because we'd lose the fractional information. We need a quantization scheme and some extra metadata to reconstruct the values.

So the compression will not be exactly 2×.

But even after including that metadata, we can significantly reduce the storage requirement.

The main design decision was to compress only the CPU-resident representation, while leaving the GPU representation unchanged.

text
Normal HiCache:

    GPU L1 (BF16)
          ↕
    CPU L2 (BF16)


HiQCache:

    GPU L1 (BF16)
          │
       Quantize
          ▼
    CPU L2 (INT8)
          │
      Dequantize
          ▼
    GPU L1 (BF16)

This gives us an important separation. The attention kernels still consume their normal BF16 KV tensors. The compression and decompression happen only when transferring KV between the two tiers.

I wasn't trying to redesign how attention works or how HiCache decides which prefixes to retain. I only wanted to change how the L2 cache stores those prefixes.

Building the INT8 codec

The first part was implementing the encoder and decoder.

The encoder takes BF16 KV values and produces a compact representation. The decoder takes that representation and reconstructs approximate BF16 values.

Since KV values are floating-point numbers, we need a way to represent them using the limited integer range of INT8.

I used symmetric, per-head INT8 quantization.

Let's understand that with a small example.

Suppose a KV head contains these four values:

text
[1.0, 0.5, -0.2, 0.0]

INT8 cannot directly represent 0.5 or -0.2. So we first find the maximum absolute value in the head and calculate a scale:

text
scale = max(abs(values)) / 127
      = 1.0 / 127
      ≈ 0.007874

We can think of the scale as the real-number value represented by one quantization step.

Next, divide each value by that scale and round it to the nearest integer:

text
Original BF16       INT8 code

  1.0       ─────►     127
  0.5       ─────►      64
 -0.2       ─────►     -25
  0.0       ─────►       0

Now we store those integers together with the scale.

During restoration, we multiply each integer by the scale:

text
Reconstructed value = INT8 code × scale

The recovered values will be close to the originals, but not necessarily identical. Quantization is lossy.

For the actual implementation, I calculate one scale independently for each KV head, for each token and layer, separately for keys and values.

Why per-head rather than one common scale for the whole row?

Because different heads can have different value ranges. A head with values mostly around 0.1 shouldn't have to share the same quantization step as another head containing values around 100.

Per-head scaling gives each head its own dynamic range, with relatively little metadata overhead.

The core operations are:

python
amax = x.abs().amax(dim=-1).clamp_min(AMAX_FLOOR)
scale = amax / 127

q = (x.float() / scale.float().unsqueeze(-1)).round()
q = q.clamp(-127, 127).to(torch.int8)

# Dequantization
x_restored = q.to(torch.bfloat16) * scale.unsqueeze(-1)

There are a few details worth mentioning.

The quotient is formed in FP32 because BF16 division would introduce avoidable rounding error when calculating the integer codes. I also explicitly clamp values to [-127, 127] because rounding the scale to BF16 can make the computed quotient slightly exceed the intended range.

And what happens when an entire KV head contains zeros?

Without protection, its scale would be zero, leading to division by zero. So I introduced a small nonzero floor for the maximum absolute value. An all-zero head still stores zero integer codes and restores to exactly zero.

These are small details, but they're important when the codec is going to sit inside an inference runtime.

The memory layout: why 1,152 bytes?

The next question was how to store the quantized values.

For Qwen3-8B at TP=1, one token has 8 KV heads, each with 128 dimensions.

For a single K or V row:

text
INT8 values = 8 × 128 × 1 byte
            = 1,024 bytes

BF16 scales = 8 × 2 bytes
            = 16 bytes

Raw encoded row = 1,040 bytes

At this point, 1,040 bytes looks like the natural record size.

But I couldn't use it directly with SGLang's existing CUDA HiCache mover.

The mover configuration used by this path requires the transferred element size to be a multiple of 128 bytes. Internally, it uses 32-thread copy groups operating on 128-byte groups.

And 1,040 is not divisible by 128.

The next valid multiple is 1,152.

So the final record becomes:

ComponentBytes
INT8 payload1,024
8 BF16 scales16
Alignment padding112
Total1,152

That gives us:

text
1,152 = 9 × 128

The 112-byte padding doesn't contain useful KV data. It exists to keep the record compatible with the existing mover.

This was a deliberate trade-off. I could try to eliminate the padding by changing the transfer implementation, but that would introduce a much larger integration surface for a relatively small storage saving.

Now let's calculate what this means for the complete KV cache.

text
Original BF16:
    2 × 36 × 2,048
    = 147,456 bytes/token

HiQCache INT8:
    2 × 36 × 1,152
    = 82,944 bytes/token

So the reduction is:

text
1 - (82,944 / 147,456)
= 43.75%

And the relative capacity becomes:

text
147,456 / 82,944
= 1.7778×

In other words, the same physical CPU memory can hold approximately 1.78× as many KV tokens.

The record layout determines this result. The next question was whether I could integrate that representation into SGLang without rewriting the existing caching system.

Integrating the codec into SGLang

SGLang's HiCache already handles the difficult parts of hierarchical caching: prefix matching, cache allocation, eviction, backup, restoration, and scheduling.

I didn't want to reimplement any of that.

Instead, I traced the path that transfers KV between the device pool and the host pool. The interesting part was that the L1–L2 traffic goes through the host-pool interface.

That gave me two places to intercept the transfers:

python
backup_from_device_all_layer(...)
load_to_device_per_layer(...)

By implementing a new host-pool class, MHATokenToKVPoolHostINT8, I could override these methods and introduce encoding and decoding at the transfer boundary.

The integration uses three main components:

ComponentResponsibility
int8_codec.pyQuantization, packing, and reconstruction
int8_staging.pyGPU staging buffers and pointer tables
mha_int8.pyCompressed host pool and transfer implementation

The rest of SGLang still sees a host pool with logical KV slots. It doesn't need to understand the encoded values stored inside those slots.

One particularly useful detail is get_size_per_token().

HiCache uses the host pool's reported bytes per token to determine how many tokens fit into the configured memory budget.

So the new pool reports:

python
def get_size_per_token(self):
    return 2 * self.layer_num * 1152

Instead of reporting the original BF16 token size, it reports the encoded size.

That connects the smaller representation directly to increased logical L2 capacity, without modifying the scheduler or radix-cache policy.

However, changing the host representation exposed another problem.

The problem with SGLang's CUDA mover

My first thought was simple: quantize the GPU KV and transfer the smaller representation to CPU memory.

But the existing HiCache mover is designed to copy bytes, not perform quantization.

It accepts separate source and destination strides, but a single transfer element size.

That means it cannot directly transform a 2,048-byte BF16 row into a 1,152-byte encoded row during the transfer.

So I introduced GPU staging buffers.

The idea is to finish the compression on the GPU first. Once the compressed bytes are ready, the existing CUDA mover can copy them normally.

GPU to CPU: backup

text
GPU L1
BF16 KV
   │
   ▼
Gather selected KV rows
   │
   ▼
Quantize + pack
   │
   ▼
GPU staging buffer
[1,152-byte records]
   │
   ▼
Existing CUDA byte mover
   │
   ▼
CPU L2
Compressed INT8 records

For backup, the implementation encodes the selected K and V rows across the model's layers and then issues an all-layer transfer using the existing mover.

CPU to GPU: restoration

Restoration happens in the opposite direction:

text
CPU L2
Compressed INT8 records
   │
   ▼
Existing CUDA byte mover
   │
   ▼
GPU staging buffer
   │
   ▼
Unpack + dequantize
   │
   ▼
Write BF16 KV into L1
   │
   ▼
Attention

Here the compressed rows are copied back to the GPU staging buffer, decoded, and written into the regular device KV pool.

I used separate persistent staging buffers for backup and restoration because the two directions can have transfers in flight concurrently.

Sharing the same buffer would create the possibility of one transfer overwriting records that another transfer hadn't finished consuming.

The buffers also grow when necessary, rather than being freshly allocated for every transfer.

This added some GPU memory overhead, but it allowed me to keep using the existing HiCache CUDA mover.

Making asynchronous restoration safe

There was another correctness issue I had to account for.

CUDA operations are asynchronous. The CPU can enqueue a transfer and continue without waiting for the GPU to finish executing it.

Now imagine restoring a cached layer:

text
1. Copy compressed KV to GPU
2. Dequantize KV
3. Write reconstructed KV into L1
4. Attention reads that KV

What happens if attention starts before step 3 finishes?

It could read stale or partially written KV values. And that type of bug might not cause a crash. It might simply produce incorrect tokens.

SGLang already has synchronization mechanisms for its transfer streams, so I reused those rather than introducing a separate synchronization system.

The essential requirement is that the compressed copy, dequantization, and scatter into the device KV pool are all enqueued on the H2D transfer stream before the layer is marked ready.

text
H2D transfer stream:

Copy encoded records
        ↓
Dequantize
        ↓
Write BF16 KV
        ↓
Record layer-ready event

             │
             ▼

Model forward stream:

Wait for layer-ready event
        ↓
Consume restored KV

CUDA preserves ordering within a stream, and the completion event establishes the dependency between the transfer stream and the model's forward stream.

This means the model cannot consume the restored KV before the required work completes.

I also had to account for the lifetime of staging allocations when buffers grow. Replacing a GPU buffer while another stream still uses it can lead to memory reuse before the queued work finishes.

These are the kinds of details that don't appear in the quantization formula, but matter just as much for a correct implementation.

Measuring whether it actually helps

Once the compressed path was working, I wanted to verify two separate things.

First, does the new representation really save the memory predicted by the byte arithmetic?

Second, does that extra capacity translate into more useful cache reuse?

I ran the experiments with Qwen3-8B at TP=1 on an NVIDIA RTX A6000 48 GB GPU.

The workload used 32 shared-prefix groups, approximately 2,048 tokens per group, 128 requests, and concurrency 8. I capped the GPU KV pool at 16,384 tokens so that L2 capacity would matter.

I used two experiments, each holding a different quantity constant.

Experiment A: same token capacity

Here I fixed both configurations to the same logical L2 capacity: 54,254 tokens.

The question was: how much CPU memory does each representation require to store the same number of tokens?

MetricBF16 HiCacheHiQCache
L2 token capacity54,25454,254
Host bytes/token147,45682,944
Host memory8.00 GB4.50 GB
L2 evictions00

So the same cache capacity fits into 56.25% of the original memory.

This is the most direct test of the storage-efficiency claim because the two configurations have identical token capacity.

One implementation detail: I used --hicache-ratio rather than --hicache-size to configure this experiment. The latter takes integer gigabyte values, while the ratio allows more precise control of the logical token count. I also checked the actual capacity reported by the running host pool rather than trusting only the configuration arithmetic.

Experiment B: same CPU memory budget

Now let's reverse the experiment.

Instead of holding token capacity constant, give both configurations the same 8 GB CPU memory budget.

The question becomes: does the smaller representation allow us to retain more useful KV, and does that improve prefix reuse?

MetricBF16 HiCacheHiQCache
CPU L2 budget8 GB8 GB
L2 capacity54,25496,451
Cache-hit rate, representative run56.1%72.7%
Tokens restored from L2126,997167,899
Mean TTFT, representative run1,244 ms922 ms
Output throughput, representative run105.1 tok/s127.8 tok/s

This was the behavior I hoped to see.

The larger L2 cache retained more of the reusable working set, which meant more requests could restore existing KV instead of recomputing their prefixes.

But I didn't want to draw broad latency conclusions from one favorable run.

Across four paired fixed-budget measurements, the cache-hit-rate improvement was relatively consistent:

text
BF16 HiCache:   56.8% – 57.6%
HiQCache INT8:  72.7% – 73.4%

The latency results were more variable.

Three of the four paired measurements showed substantial TTFT improvement, while one was near parity or slightly worse.

So I wouldn't describe HiQCache as a universally faster inference system. The more defensible result is that under this workload, the larger effective L2 capacity produced consistently more cache reuse.

And that distinction matters.

The extra cache capacity is a direct consequence of the record size. The serving improvement depends on whether the workload actually benefits from retaining those extra prefixes.

For example, if every reusable prefix already fits in BF16 L2, compression may not increase the cache-hit rate at all. In that case, the extra encoding and decoding operations could make HiQCache slower.

What about quantization error?

There is one important trade-off that memory benchmarks cannot answer.

INT8 quantization is lossy.

When HiQCache restores a compressed row, it produces BF16 values, but those values are not necessarily identical to the original BF16 KV.

So how much error are we introducing?

I tested the codec against an elementwise numerical error bound that accounts for INT8 rounding, storing scales in BF16, and rounding the reconstructed values to BF16.

The bound is:

text
|x̂ - x| ≤ (0.5 + 2⁻⁸) × s + 2⁻⁸ × |x̂|

The numerical tests verified this bound across CPU, CUDA, and Apple MPS, with no violations in the tested inputs.

But there is an important distinction between numerical correctness and model-level quality.

Even when individual KV values are reconstructed with small errors, those errors can propagate through attention and subsequent transformer layers. Sometimes small changes in the attention output can change the final token selection.

I built a quality-comparison harness to measure token agreement and logprob differences between BF16 and INT8 restoration.

However, the recorded quality experiment did not actually trigger L2 restoration in either configuration. So those requests didn't exercise the codec on the read path, and the result couldn't establish the quality impact.

Rather than report that comparison as valid, I marked the result invalid and added a check requiring actual L2 restore activity.

So the current evidence establishes that compressed restoration works and that the codec satisfies the tested numerical bound. It does not yet establish that generation quality is unchanged.

That is still an important piece of validation to complete.

What I deliberately kept out of V1

The current implementation is intentionally narrow.

It targets Qwen3-8B with TP=1, 8 local KV heads, 128 dimensions per head, page size 1, and the layer_first host layout.

I haven't generalized the compressed record to tensor-parallel configurations, MLA, larger page sizes, or L3 storage.

For example, increasing tensor parallelism changes the number of local KV heads. That would change the payload size, scale metadata, padding, and transfer geometry. The fixed 1,152-byte layout is no longer appropriate without redesign.

Rather than silently support a configuration that could produce incorrect KV, the host pool rejects unsupported geometries at initialization.

I also kept the codec implemented with PyTorch GPU operations rather than immediately writing fused Triton kernels.

A fused implementation might reduce kernel launches and intermediate operations, but I first need reliable per-phase profiling to know whether encoding and decoding are actually the bottleneck.

There is no reason to introduce another specialized kernel solely because it sounds faster.

What I learned from building HiQCache

The quantization itself was probably the simplest part of the project.

The more interesting work started when I tried to make that smaller representation fit into an existing inference runtime.

A 1,040-byte record looked reasonable until the CUDA mover required 128-byte alignment. Directly copying BF16 into compressed storage looked straightforward until I traced the mover's single-element-width contract. And restoring a compressed row looked like a normal copy-and-decode operation until I had to account for transfer streams, completion events, and GPU allocation lifetimes.

Each of these constraints shaped the final design.

The result is a relatively small change to SGLang's host-cache representation, without requiring changes to the attention kernels, radix cache, scheduler, or existing HiCache CUDA mover.

The most important outcome is that I could store the same logical KV cache in 43.75% less CPU memory, or retain 1.78× as many tokens under the same memory budget. In the measured shared-prefix workload, that extra capacity also produced more cache reuse.

There is still work to do, particularly deterministic end-to-end quality comparisons and detailed codec profiling. But that is also part of what made the project useful for me: understanding not just how to implement an optimization, but what needs to be measured before claiming that it is actually better.