KV-Caching, Batching, and Memory Complications : Scaling Kshipra on CPU

September 23, 2026

The first Kshipra post ended on an unambiguous problem: the baseline generator recomputes attention over the entire sequence on every single step, creating O(N²) computational waste. Fixing that meant tackling two things — Key-Value (KV) caching, to stop redoing work, and static batching, to squeeze more out of the CPU while I was at it.

Along the way I ran into a subtle left-padding bug that quietly corrupted position embeddings, and a memory result that looked backwards until I understood what psutil was actually measuring.

1. KV-Caching: Converting O(N²) to O(1)

During autoregressive generation, attention for a new token is computed against every historical key and value. Since past tokens never change, their key/value activations don't need to be recomputed on every step — they just need to be cached and reused.

That splits generation into two phases:

  • Prefill (step 0): process the full prompt in one forward pass, compute K and V for every prompt token, and store them as past_key_values.
  • Decode (every step after that): pass in only the single newest token (shape (B, 1)), alongside the cached past_key_values. PyTorch appends the new key/value states to the cache automatically.
# Decoding pass in the KV-Cache generator
outputs = self.model(
    input_ids=next_token_id,        # Shape: (B, 1)
    past_key_values=past_key_values,
    use_cache=True,
)
past_key_values = outputs.past_key_values  # Persist state for the next step

Result: per-token decode latency dropped from a rising linear curve to a strictly flat O(1) line — around 0.058 s/token — pushing throughput from 5.52 tok/s to 16.45 tok/s.

2. Static Batching & the Left-Padding Position Bug

To push CPU throughput further, I added static sequence batching (B=4) in engine/batching.py. Grouping 4 prompts into a single tensor shifts the workload from memory-bandwidth-bound matrix-vector products to compute-efficient matrix-matrix operations — a much better fit for a CPU.

The Left-Padding Requirement

In decoder-only architectures like GPT-2 or Llama, prompts in a batch need to be left-padded, so the last real (non-pad) token in every sequence lines up in the same rightmost tensor column when generation begins.

GPT-2 doesn't ship an official padding token, so I pointed it at EOS instead:

self.tokenizer.padding_side = "left"
if self.tokenizer.pad_token is None:
    self.tokenizer.pad_token = self.tokenizer.eos_token

The Bug: Repeating Token Loops

The first time I ran batched generation, the longest prompt in the batch came out clean. Every shorter prompt collapsed into a repeating loop of garbage — things like artificial AI AI AI AI....

Root cause: by default, transformer models compute position IDs as a plain range, [0, 1, 2, ..., T_max - 1]. With left-padding, a short prompt like ["[PAD]", "[PAD]", "The", "future"] was getting position IDs [0, 1, 2, 3] — telling the model that the real first token, "The", sits at absolute position 2, not position 0. That single off-by-N error corrupted the absolute positional embeddings for every step of generation that followed.

The Fix: Explicit Position ID Indexing

The fix is to derive position_ids from the cumulative sum of the attention_mask, so padding tokens get excluded from the count entirely:

# 1. Prefill phase: offset position_ids so the first real token gets index 0
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 0)
 
outputs = self.model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    position_ids=position_ids,
    use_cache=True,
)
 
# 2. Decode phase: compute the exact index for the new single token
position_ids = attention_mask.long().sum(dim=-1, keepdim=True) - 1
 
outputs = self.model(
    input_ids=next_token_id,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_values=past_key_values,
    use_cache=True,
)

With this fix, batched output matched single-request, sequential generation character-for-character.

3. The Memory Paradox: Why Did Batched RAM Go Down?

Once the bug was fixed, the benchmark numbers threw me a second surprise:

  • Sequential KV-Cache: 80 total forward passes (4 × 20), 16.45 tok/s, 1208.04 MB peak RAM
  • Batched KV-Cache (B=4): 20 total forward passes, 35.17 tok/s, 1147.02 MB peak RAM

Storing four KV-caches at once, somehow, used about 61 MB less peak RAM than processing the same four requests one at a time. That's backwards from what I expected, and it took some digging to understand why.

RSS isn't tensor memory. psutil reports Resident Set Size — the total physical RAM the OS has handed to the Python process — not the size of the actively-used PyTorch tensors at any given instant. The two numbers are related, but they're not the same thing.

Allocator fragmentation is the real culprit. Sequential execution needed 80 total forward passes to finish the same workload. Each pass allocates and frees a batch of temporary activation tensors. Repeating that 80 times gives PyTorch's underlying C++ memory allocator plenty of chances to fragment: it holds onto freed blocks instead of returning them to the OS, on the assumption that a same-sized allocation is probably coming right back. That habit is what drove peak RSS up to 1208 MB.

Fewer passes, less churn. Batched execution does the exact same amount of work, but in only 20 forward passes instead of 80 — a 75% cut in temporary allocate/free cycles. Less churn means less fragmentation, which is what kept peak RSS at 1147 MB.

The KV-cache itself barely moves the needle here. For GPT-2 (12 layers, 12 heads, 64-dim heads, FP32), a 30-token KV-cache at B=1 is only about 0.66 MB; at B=4 it's about 2.65 MB. That ~2 MB difference is nowhere near enough to explain a 61 MB gap on its own — it's completely dwarfed by the allocator savings from running fewer forward loops.

Benchmark Summary

  • Throughput speedup: 2.14× net gain — 35.17 tok/s batched vs. 16.45 tok/s sequential.
  • Latency profile: a flat ~0.11 s inter-token step latency, holding steady across every generation step. Between KV-caching and batching, Kshipra went from 5.52 tok/s at baseline to 35.17 tok/s — over 6× faster — without touching a GPU. Next up on the list: dynamic batching and INT8 quantization, both aimed at squeezing more throughput out of the same 8GB ceiling.