inkling: ring-buffer KV cache for sliding-window layers (~10x less KV memory at long context)

[dpanelli]

TL;DR

Sliding-window layers never read KV rows older than window, but their cache was sized for the whole context. This PR keeps only the last window rows per sliding layer, addressed as a t % window ring. Output is bit-identical, verified token-for-token against the HF reference with the ring wrapping. On the real 975B config (55 of 66 layers sliding) total KV memory drops from 30.3 GB to 3.2 GB at 32k context and from 121 GB to 11.4 GB at 128k; the numbers are arithmetic from the checkpoint's config over the engine's f32 cache, and the engine's own [kv] report validates the same formula through the shipped allocation path. One file, no new knobs, no semantic change.

What this does

Sliding-window layers in inkling.c only ever attend to the last window positions: attention() clamps the scan at t0 = qpos - window + 1. The cache allocation ignored that. kv_alloc sized every layer's K/V for the full max_t, so on the real Inkling shapes (66 layers, 55 of them sliding, window 512) a sliding layer at 32k context held 32,768 rows of cache and could ever read at most 512 of them. Rows older than the window are written once and then dead.

This PR makes each sliding layer's cache a ring of window rows, addressed at t % window. Global layers keep the full cache. Output is bit-identical.

How it works

All in c/inkling.c:

• kv_ring_rows(c, li, max_t) returns window for sliding layers when 0 < window < max_t, else max_t. Both kv_alloc and attention derive sizes from it, so allocation and addressing cannot diverge. This also covers the kv_alloc reuse path, where win comes from the allocated m->max_t rather than the requested one.
• Scoring reads rows of the current batch (t >= pos0) from the k/vv scratch buffers and older history from the ring. The scratch holds exactly the bytes the cache would hold (post-sconv, post-rmsnorm), so the arithmetic is unchanged. This is needed because the append moved after the scoring loop: with a ring, appending the whole batch up front can overwrite history rows that earlier queries in the same batch still need (any prefill with S > window).
• The append skips rows the same batch would immediately overwrite (s0 = max(0, S - win)). Why that is safe: skipped rows are positions < end - window, and no later query attends earlier than end - window + 1.

Two smaller additions: a fail-fast bound check at the top of attention() (an over-run used to be an OOB write; with the ring it would have wrapped silently and produced wrong output instead of crashing), and a one-line [kv] allocation report on stderr.

Why

KV memory is the context ceiling for this engine. After int4-gs64 conversion the dense weights of the 975B fit a 25 GB box, but the f32 KV cache at 32k context was 30 GB, larger than the weights, and 92% of it was unreachable by construction. With the 5:1 sliding:global layout the fix applies to 55 of 66 layers.

Tests done

• Bit-exactness against the HF oracle (tools/make_tiny_inkling.py, transformers 5.14: window 8, 7 of 8 layers sliding, 12-token prompt + 24 generated, so the ring wraps 4+ times): teacher-forced argmax 36/36, greedy generation 24/24, exit 0. Token-identical before and after the change. The teacher-forced pass (S=36 in one batch, so S > window) exercises the scratch-read path; greedy decode exercises ring reads after wraparound.
• make test-c passes; -Wall -Wextra clean.
• Memory-safety and concurrency review: int64 promotion on every cache offset; the append sits outside the OpenMP parallel region (ring, k, vv are read-only during scoring); state_reset not clearing K/V stays correct because a fresh request's first batch performs zero ring reads, and every later read hits a slot written earlier in the same request. The index mapping was additionally checked by exhaustive simulation across prefill, decode, window >= max_t, and reuse over a dirty ring.

Expected memory savings

For the real Inkling config (66 layers: 55 sliding with 16 KV heads, 11 global with 8, head_dim 128, f32 cache):

context	full KV	ring KV	factor
8k	7.6 GB	1.1 GB	6.8x
32k	30.3 GB	3.2 GB	9.5x
128k	121 GB	11.4 GB	10.6x

Per-token KV growth drops from 968 KB to 88 KB, since only the 11 global layers keep growing. The 55 sliding layers cost a flat 440 MB total at any context length. At 128k the full cache alone (121 GB) was bigger than the RAM budget of a CPU box that fits the whole engine; with the ring it is 11.4 GB.

On the tiny oracle model the self-report reads [kv] 0.0 MiB (ring buffers on sliding layers; full cache would be 0.2 MiB).

Out of scope

• Global layers are untouched. Shrinking those is a different problem (fewer full-attention layers, or cache quantization like the KV8/TQ work upstream).
• The t % win in the scoring loop could be split into two contiguous runs to shed the division and the branch. Left simple: output is identical either way, and for decode the cost is noise next to the head_dim-128 dot products. Worth revisiting only if a profile says so.

🤖 Generated with Claude Code

https://claude.ai/code/session_01828WtkfFqicywaZLZu1139

[JustVugg]

This is a good change and the reasoning is right: a sliding layer at 32k holding 32,768 rows to read at most 512 of them is pure waste, and 30.3 GB → 3.2 GB is worth having. One file, no new knobs, bit-identical output — that is the shape a change like this should have.

It conflicts with dev now, and the conflict is in the one place where a mechanical resolution would be wrong. GitHub may still show this as mergeable; it is recomputing. Locally it is a real conflict in kv_alloc.

dev gained grow-with-copy in #786, for KV prefix reuse:

float **oldK = m->K, **oldV = m->V;
int old_max = m->max_t;
int keep = (m->K && m->kv_len > 0 && m->kv_len <= max_t) ? m->kv_len : 0;
...
for (int h = 0; h < kv && keep; h++)
    memcpy(m->K[i] + (int64_t)h * max_t  * hd,
           oldK[i] + (int64_t)h * old_max * hd, (size_t)keep * hd * sizeof(float));

That loop assumes position t lives at row t. Your ring breaks that assumption for sliding layers, where position t lives at row t % win. Copy the first keep rows linearly out of a ring that has already wrapped and you get a silently rotated cache — bit-identical output on short prompts, wrong attention as soon as a conversation runs past window and reuse engages. It would not fail loudly.

The good news is that the ring makes the sliding case simpler, not harder:

• Sliding layers. rows is window, which does not depend on max_t. The buffer never needs to grow, so there is nothing to re-lay-out — copy the ring wholesale (rows entries, not keep) and t % win stays valid, or better, skip the free/realloc for those layers entirely.
• Global layers. rows is still max_t, the stride still changes, and the existing linear copy is exactly right.

So the resolution is a branch on kv_ring_rows(c, i, max_t) < max_t, not a merge of the two hunks.

One more thing to keep whole while you are in there: kv_alloc ends with

if (kv_prefix_grow(&m->kvp, max_t, keep)) m->kv_len = keep;
else                                      m->kv_len = 0;

The version in this PR predates that line. It has to survive, or prefix reuse stops on Inkling and nothing says so.

I would rather you resolve this than have me guess at it — it is your ring and my grow-with-copy colliding, and the failure mode is quiet. Ask if any of the above is unclear; #786 has the reasoning behind the copy.

Your CI had never run, incidentally — it sat in action_required, which GitHub applies to a first contribution from a fork and surfaces nowhere on the PR page. I approved it: 13 checks, no failures. So the change itself is sound; it is only the collision that needs a hand.

[dpanelli]

Rebased on dev and fixed conflicts