The previous Kshipra post ended with batching and quantization as next steps under an 8GB ceiling. Taking that further meant moving from GPT-2 to a Llama-style model with grouped-query attention (GQA), and the old design did not survive the move. Growing KV tensors on the fly, padding batches and recomputing every prompt do not fit a 4 GiB peak-RSS budget on a CPU.
So I rebuilt the engine around one question: where do the bytes go? Decoding one token streams the model's weights through the memory bus, so on a CPU, memory bandwidth and residency matter more than FLOPS.
Along the way I hit a quantization bug that broke the error bound, a cache-eviction bug that freed half the requested blocks, and a memory test that measured its own temporary tensors.
1. Block-Arena KV Cache
The old cache allocated KV tensors as sequences grew. That fragments the allocator and makes peak memory hard to predict.
The new cache reserves KV memory once, as an arena of fixed-size blocks. A LIFO free-list hands blocks out and takes them back:
# simplified
free = list(range(num_blocks)) # pop() returns the highest block ID first
block = free.pop() # allocate
free.append(block) # free: the next allocation reuses this blockLIFO reuse means the most recently freed block is the next one used, so it is the most likely to still be in cache. No allocation happens per token.
Sizing is GQA-aware. KV bytes per token = 2 × layers × KV heads × head dim × bytes per element. For a Llama-3-8B-shaped config (32 layers, 8 KV heads, head dim 128, FP16) that is 128 KiB per token. With 32 full KV heads it would be 512 KiB. At 4,096 tokens that is 512 MiB versus 2 GiB, and under a 4 GiB ceiling that gap decides what fits.
2. INT4 Quantization and the Scale Bug
Weights are quantized to INT4 per group of 128, with two nibbles packed per byte. Each group stores one FP16 scale, so the cost is 4 + 16/128 = 4.125 bits per weight, about 26% of FP16.
Rounding to the nearest step guarantees each weight is reconstructed within half a step. That guarantee holds only if encoding and decoding use the same scale.
The Bug: Error Bound Violation
The round-trip tests failed by small margins.
Root cause: weights were quantized with the FP32 scale and dequantized with its FP16 rounding.
q = round(w / s_fp32) # encode
w_hat = q * s_fp16 # decode, s_fp16 = fp16(s_fp32)
# bound to hold: |w - w_hat| <= s / 2FP16 rounds a scale by up to 2^-11 relative, so the error grows by up to |q| × 2^-11 × s beyond s/2. For INT8 (|q| up to 127) that is about 0.06 s. The bound is tight at s/2, so it fails by exactly this kind of margin.
The Fix: One Scale on Both Sides
scales_for_quant = scales_fp16.to(torch.float32)
q = round(w / scales_for_quant)Round the scale to FP16 once and use that value for encode and decode. The INT8 and INT4 round-trip bound tests now guard it.
3. Radix-Tree Prefix Cache and the Eviction Bug
Many requests share a prefix, such as a system prompt. Recomputing it per request wastes the most expensive step.
KV blocks live in a radix tree keyed by token blocks. A new request walks the tree, reuses every matching block and computes only the rest. Reference counts keep shared blocks alive. If a request diverges inside a shared block, that block is copied before it is written (copy-on-write). Under memory pressure the cache evicts the least recently used leaf. Nodes use flat __slots__ and a token arena instead of nested Python objects, to avoid pointer-chasing.
The Bug: Stale Leaf Snapshot
evict(2) freed 1 block.
Root cause: evict listed the leaves once, then looped over that list. Evicting a leaf can turn its parent into a leaf, but the parent was not in the snapshot, so it was never considered.
The Fix: Re-Scan Every Iteration
# simplified
while freed < target:
leaf = lru_leaf(tree) # re-scan the tree each iteration
if leaf is None:
break # nothing evictable left
freed += evict(leaf)The cost is O(nodes) per eviction. That is fine at current sizes. If profiling flags it, an LRU-ordered set of evictable leaves, updated when a parent loses its last child, removes it.
4. Chunked Prefill
Prefilling a long prompt in one pass stalled decoding for about 350 s on my machine. Every active stream froze.
Chunked prefill splits the prompt into chunks and interleaves them with decode steps. Attention over a chunk runs an online softmax across KV blocks that are read in place from the arena. That skips a gather copy of about 256 MiB.
# per query row, simplified
m, l, o = float("-inf"), 0.0, 0.0
for k_blk, v_blk in kv_blocks:
s = q @ k_blk.T * scale
m_new = max(m, s.max())
p = exp(s - m_new)
l = l * exp(m - m_new) + p.sum()
o = o * exp(m - m_new) + p @ v_blk
m = m_new
out = o / lTwo details matter for correctness:
- RoPE uses absolute positions. A chunk starting after P cached tokens gets positions P onward, not 0 onward.
- The causal mask covers the chunk. Each query attends to the cached prefix and to earlier tokens in its own chunk only.
The merge is mathematically identical to dense attention, and the tests confirm chunked output matches dense masking within floating-point tolerance.
5. Continuous Batching and Admission Control
Static batching, from the last post, pads every request to one length and waits for the slowest. Continuous batching schedules per iteration. Each step mixes prefill chunks and decode tokens under a token budget, Sarathi-style.
The risk is memory deadlock: admit too many sequences and none can finish. Admission is strict:
free + evictable >= active_reserved + new_needIf memory is still short, a running sequence can be preempted. Recomputing its KV cache on a CPU is expensive (one recompute cost about 178 s), so preemption swaps its blocks out of the arena and restores them later.
Status: implemented, not yet validated under load.
6. Measuring Memory Correctly
The last post showed that psutil RSS reflects allocator behavior, not tensor memory. Now every test runs in a fresh spawn subprocess and reads VmHWM from /proc/self/status, the kernel's peak RSS.
Two problems cost me more time than any engine bug.
The CUDA Wheel
A default pip install torch pulled in about 2 GB of CUDA libraries. The CPU wheel cut baseline RSS from about 400 MB to about 185 MB.
The Test That Measured Itself
A 33 MiB INT4 tensor showed a growth of about 290 MiB.
Root cause: the test forced pages into RAM with .sum(). On a uint8 tensor, .sum() accumulates in int64, which allocates a temporary eight times the size, about 268 MiB.
Fix: read one byte per 4 KiB page.
t.view(torch.uint8).view(-1)[::4096].sum()Summary
- Measured so far: baseline RSS of about 185 MB on the CPU wheel, INT4 at 4.125 bits per weight, a monolithic prefill stall of about 350 s, and a recompute cost of about 178 s.
- Verified: greedy decoding matches HuggingFace
generatetoken for token on the full-precision path, chunked attention matches dense masking, and eviction cascades correctly under memory pressure. - Not yet measured: end-to-end throughput and latency for the new engine, continuous batching under load, and the 4 GiB ceiling across a full run.
Next is validating the scheduler, then load-testing it with asynchronous Poisson arrivals and around 100 concurrent users.