[RFC/Experiment] Colibri Native Residency Engine: cost-aware expert caching and paged KV prefixes

[terrizoaguimor]

@JustVugg, I would like your opinion on this direction before anyone writes implementation code.

I reviewed the current cache, storage, and KV mechanisms across the five model families available on dev: GLM-5.2, Kimi K3, Inkling, OLMoE, and DeepSeek V4.

My conclusion is not that Colibri needs an external cache service. It needs a common native residency engine that applies Valkey-like bounded-memory and eviction discipline to inference-native objects.

There would be no daemon, sockets, RESP protocol, serialization layer, RDB, or AOF. The system would live inside the C runtime and operate directly on expert bundles, prepared tensor layouts, and KV blocks.

The objective would not be to maximize cache hit rate. It would be:

Under fixed RAM, VRAM, and storage-bandwidth constraints, keep the objects that remove the most expected critical-path time per resident byte.

Why this may be worth investigating

Colibri already has several effective caches, but each engine implements a different subset:

Engine	Expert RAM cache	Persistent pins	Dynamic repin	VRAM tier	Prefix/KV reuse
GLM-5.2	per-layer LRU	yes	yes	CUDA and Vulkan	multi-slot, persisted KV, cross-slot adoption
Kimi K3	per-layer LRU	yes	no	Vulkan	exact prefix, one session
Inkling	per-layer LRU	yes	no	Metal zero-copy	exact prefix, one session
OLMoE	per-layer LRU	partial	no	no	incremental private chat
DeepSeek V4	lease-aware ExpertStore	yes	yes	no	exact prefix, one session

The project already has most of the required mechanisms:

• per-layer LRU caches;
• learned .coli_usage histories;
• fixed and adaptive pins;
• direct I/O and buffered I/O;
• mmap, mirrors, and multiple drives;
• PIPE and io_uring;
• CUDA/Vulkan/Metal residency;
• lease-protected expert objects in DeepSeek V4;
• persisted KV and exact prefix reuse;
• continuous batching in GLM.

The missing piece is a common mathematical and lifecycle model.

Adding more model families currently means implementing another cache structure, another admission policy, another set of counters, and another interpretation of memory capacity. Qwen and other future engines will increase that duplication unless there is a reusable seam.

This proposal is intended to reduce that friction, not add a new framework around every model.

Evidence from current measurements

Cache residency can be a major performance lever when storage is on the critical path:

• Full expert residency moved a six-GPU GLM run from approximately 2.30 to 5.77 tok/s and removed disk wait.
• Raising useful RAM on a Ryzen 5950X moved throughput from 0.77 to 1.02 tok/s.
• On an M5 Max streaming configuration, approximately 60% of wall time was felt expert-I/O wait.
• Exact KV-prefix reuse reduced a measured second-request TTFT from 17.689 to 0.569 seconds.
• Another long-prefix case reduced TTFT from approximately 50.1 to 1.7 seconds.

But the measurements also show why hit rate alone is not the correct objective:

• A workload-trained placement improved an in-sample run from 8.74 to 13.01 tok/s but regressed held-out throughput from 8.93 to 5.54 tok/s.
• Increasing prefetch queue depth has sometimes increased bytes read and reduced throughput.
• Fully resident GLM is no longer cache-bound: expert matmul, attention, memory bandwidth, and orchestration dominate.
• Bad NUMA placement has reduced effective expert bandwidth from roughly 65 GB/s to 31 GB/s.
• Metal has demonstrated cases where increasing logical hit rate reduced throughput because residency bookkeeping became more expensive.
• OLMoE has shallow global popularity but strong temporal locality, so the same pinning policy is not optimal for every model.

The policy must account for the actual critical path, topology, and object-preparation cost.

Two different cache problems

Expert residency

The key is approximately:

(model, layer, expert, format, backend layout)
    -> execution-ready expert bundle

A hit can avoid storage reads, repacking, decompression, or GPU upload.

This primarily affects decode on systems where the full routed-expert bank does not fit in RAM and VRAM.

Token-prefix residency

The key is approximately:

(model and execution identity, rendered token prefix)
    -> computed KV/recurrent state

A hit avoids prefill computation and primarily improves TTFT.

These must remain separate object classes with separate policies and metrics. An expert hit and a prefix-token hit do not save the same work.

Proposed architecture

I suggest a Colibri Native Residency Engine, or CNRE, embedded in the C runtime.

It would manage inference-native objects rather than generic byte strings:

ExpertObject
PreparedExpertObject
CompressedExpertObject
KVPrefixBlock
KVPrivateTail

Object identity

An expert object should include:

model content identity
engine family
layer
expert id
quantization format
group size
tensor geometry
packing version
backend layout version

A KV object needs additional identity:

tokenizer identity
chat-template identity
rendered prefix
position/RoPE configuration
KV representation version
adapter or LoRA identity
multimodal payload identity

Tokens alone are insufficient. Inkling audio already demonstrates this: two payloads may use identical placeholder tokens while carrying different audio frames.

Object lifecycle

Every cacheable object would follow a state machine:

ABSENT
LOADING
READY
IN_USE
EVICTABLE
DEMOTING

A READY object is published only after its entire read, validation, decompression, packing, and backend preparation have completed.

Duplicate misses should use single-flight loading:

request A -> begins loading expert X
request B -> joins the same load

Active users hold leases. An object with a live lease cannot be evicted while a CPU matmul or GPU operation still references it.

DeepSeek V4's ExpertStore already contains much of the correct lifecycle contract and could be the architectural starting point.

Memory tiers

Tier 0: VRAM execution-ready

Contains the exact representation consumed by CUDA, Vulkan, or Metal.

A tensor that still requires repacking or upload is not a complete Tier-0 hit.

Tier 1: RAM execution-ready

Contains CPU-ready expert slabs or objects prepared for efficient upload.

Pools should be separated by NUMA node, format, and geometry where necessary.

Tier 2: compressed RAM

Optional and experimental.

A compressed hit materializes into a bounded working slot before execution. This tier should exist only when:

decompression and materialization cost
<
storage miss cost avoided

The weights are already heavily quantized, so large compression ratios should not be assumed. Existing ANS results suggest a real but incremental opportunity, not a universal 2x gain.

Tier 3: mapped or direct-I/O backing

Depending on measured platform behavior:

buffered pread and page cache
direct I/O into final slabs
mmap with controlled pre-touch
multi-drive mirror or striping
io_uring

Tier 4: checkpoint storage

Expert weights are immutable and reconstructible.

Eviction should therefore be:

discard clean object

not:

write object to operating-system swap

The checkpoint is already the specialized backing store for weights. Only derived state such as KV may justify writeback or persistence.

Mathematical policy

For object i, candidate tier t, and planning horizon H:

reuse_i(H) = estimated probability of reuse within H
miss_i     = expected critical-path cost if absent
hit_i,t    = expected critical-path cost from tier t
bytes_i,t  = resident bytes in tier t
move_i,t   = amortized promotion/demotion cost

Define utility as:

utility(i,t) =
    reuse_i(H) * (miss_i - hit_i,t)
    --------------------------------
                 bytes_i,t

    - move_i,t / H

The important input is critical-path cost, not loader service time.

A 20 ms read that completely overlaps computation does not save 20 ms of visible latency when cached. Colibri already distinguishes read service from felt wait in parts of the telemetry; the residency engine should use felt wait.

Admission rather than automatic pollution

Many LRU designs are effectively:

miss -> load -> admit -> evict victim

That lets one-use experts evict objects with demonstrated reuse.

A better initial structure would be:

window
probation
protected

A demand miss can execute from a temporary working slot without automatically entering the resident cache.

Admission occurs only when:

utility(candidate) >
    utility(victim) + hysteresis

Frequency, recency, and reuse distance should be kept as separate signals:

frequency: number of uses, with decay
recency: time or accesses since last use
reuse distance: unique objects accessed between two uses

Reuse-distance distributions can estimate:

P(hit | cache capacity = k)

without requiring a predictive model or transition table in the first implementation.

Unlike a general key/value system, Colibri's object population is small enough to maintain exact counters: roughly 20,000 GLM experts and 82,000 Kimi experts.

Dynamic budgets per layer

Using the same number of cache slots for every MoE layer is not generally optimal.

For every layer l, estimate:

Benefit_l(k) =
    critical-path time avoided with k resident slots

The marginal value of another slot is:

Marginal_l(k) = Benefit_l(k + 1) - Benefit_l(k)

Allocate the next available slot or byte to the layer with the largest marginal benefit per byte.

This water-filling approach naturally gives:

• more capacity to flat-routing layers;
• less capacity to layers with a small stable hot set;
• more capacity where misses reach the critical path;
• less capacity where asynchronous loading hides the miss;
• correct accounting for models whose layer or expert sizes differ.

The same policy can support future model families without hardcoded assumptions about expert count, layer count, or tensor width.

Preventing migration churn

Colibri's measurements show that frequent repinning can regress performance through rereads, uploads, and synchronization.

A promotion should happen only when:

expected future savings >
    migration cost * safety factor

The engine should enforce:

• hysteresis;
• minimum object age;
• minimum interval between migrations;
• bytes-moved budget per token;
• maximum promotions per pass;
• separate staging-memory budget;
• measured contention cost.

Extensibility for future model families

The residency engine should not contain a growing switch statement for GLM, Kimi, Inkling, OLMoE, DeepSeek, Qwen, and every future model.

Each engine should provide a small adapter describing capabilities and object operations:

identify object
report size in each representation
load object from backing storage
prepare execution-ready representation
release representation
describe whether the object is CPU/GPU executable
report topology and tier constraints

The generic residency core should know nothing about model semantics, tensor names, or routing math.

A new model engine would choose which capabilities it supports:

CAP_EXPERT_CACHE
CAP_DIRECT_IO
CAP_MMAP
CAP_COMPRESSED_RAM
CAP_VRAM_RESIDENCY
CAP_KV_PREFIX
CAP_KV_PERSISTENCE
CAP_CONTINUOUS_BATCHING

Unsupported capabilities remain absent rather than emulated.

This would preserve Colibri's current one-C-file-per-model style while avoiding a new cache implementation per model.

The intended integration path should be incremental:

existing engine-local cache remains valid
new engine can adopt CNRE through an adapter
no model must migrate before the common policy is proven

KV-prefix extension

The second object family would be paged host KV.

Immutable prefix blocks

Split host KV or recurrent state into bounded blocks, potentially beginning with the 64-token granularity already used by parts of the CUDA KV path.

KVPrefixBlock {
    identity
    parent prefix
    token range
    host KV pages
    reference count
    recomputation cost
}

Copy-on-write tails

Conversations sharing the same system prompt or long prefix reference the same immutable blocks.

When they diverge:

shared immutable prefix
+ private mutable tail

Prefix-aware scheduling

Before copying KV between slots, route the request to the slot already owning the longest compatible prefix.

Copying becomes a fallback rather than the first solution.

KV eviction policy

KV utility should be:

probability of prefix reuse
* prefill time avoided
/ resident KV bytes

A 123 MB prefix that avoids tens of seconds of prefill can be more valuable per byte than many expert objects.

What to borrow from Valkey

Useful principles:

• strict max-memory accounting;
• selectable policies;
• LFU with decay;
• admission versus eviction;
• per-object metrics;
• protected active objects;
• separate resident and transient budgets;
• pools;
• predictable failure behavior under memory pressure.

Things not proposed:

• no external service;
• no sockets;
• no RESP;
• no generic strings or blobs;
• no RDB;
• no AOF;
• no replication;
• no writeback of clean expert weights.

What this does not solve

A residency engine will not raise the fully resident compute ceiling by itself.

When all useful experts are already in RAM or VRAM, current measurements show that the limiting terms are:

expert matmul
attention
DRAM bandwidth
kernel selection
NUMA placement
GPU/CPU overlap
orchestration

This proposal targets disk-bound and mixed-residency systems, plus TTFT from repeated prefixes. It should not be presented as a universal 2x decode optimization.

Experimental plan

I recommend approval only for Phase 0 first.

Phase 0: offline simulator

Replay real routing traces through:

• current LRU and pin policy;
• LFU and LFRU;
• segmented LRU;
• TinyLFU-like admission;
• cost-aware admission;
• uniform versus dynamic per-layer budgets;
• optional compressed-RAM capacity.

Simulator inputs:

route trace
expert sizes
tier
measured service time
measured felt wait
physical bytes
NUMA location
page-cache state
migration cost

Outputs:

hit rate by tier
bytes avoided
critical-path time predicted
churn
admissions and rejected admissions
per-layer capacity
migration bytes

Advance only if held-out workloads predict:

at least 10% less felt wait
and no principal workload regresses by more than 3%

Phase 1: telemetry normalization

Give every engine comparable measurements without changing decisions:

tier hit
physical miss
logical bytes
felt wait
service time
admission
rejection
eviction
migration bytes
reuse distance
single-flight joins

Phase 2: opt-in GLM RAM policy

Add one experimental policy:

COLI_CACHE_POLICY=cost

No new prefetch, no routing change, and no output change.

Compare against the existing policy under exactly the same memory budget.

Phase 3: second-engine adapter

Adopt the common interface in one structurally different engine, likely Kimi or DeepSeek V4, to prove the abstraction is not GLM-specific.

Phase 4: paged KV-prefix index

Begin in GLM because it already has multiple slots, persisted KV, cross-slot adoption, and continuous batching.

Phase 5: compressed RAM tier

Only if the simulator and decompression measurements show a positive disk-bound operating region.

Phase 6: shared-memory multiprocess mode

Only if Colibri later runs multiple processes for the same model and measurements show meaningful duplicate anonymous residency. This is not part of the default architecture.

Acceptance criteria

Expert residency

On at least two disk-bound systems:

>=5% median tok/s improvement
or >=10% p99 token-latency improvement

And:

<=2% overhead when fully resident
byte-identical output for placement-only policies
strict RAM/VRAM budget compliance
no increase in unexplained physical I/O

KV prefixes

For workloads with reusable prefixes:

>=2x TTFT improvement

And:

no incompatible-state reuse
bounded memory
correct request/session isolation

Method

• exact commit;
• fixed model or container;
• frozen .coli_usage;
• fixed DRAFT and sampling;
• documented page-cache state;
• N>=5 interleaved runs;
• medians and spread;
• held-out workloads;
• negative results published.

Contribution commitment

If Phase 0 demonstrates a robust held-out benefit and the maintainers approve the direction, I would be willing to dedicate myself to implementing and integrating it incrementally across Colibri.

This would be an upstream-first contribution with no strings attached:

• no external service, product, or commercial dependency;
• no private implementation that the repository depends on;
• no requirement for special ownership or control;
• no new dependency imposed on users;
• no expectation that the maintainers accept a large rewrite at once.

The design, simulator, experimental methodology, results, implementation notes, and user-facing documentation would all be developed directly in this repository, under the project's existing license, review process, and maintainer governance.

I would simply do the work. The resulting architecture and documentation would become part of Colibri and remain available to the project and its community.

Integration would remain incremental and reviewable:

1. simulator and mathematical model;
2. normalized telemetry;
3. one opt-in policy in one engine;
4. held-out validation;
5. a second engine proving the abstraction is reusable;
6. broader adoption only after each stage demonstrates value.

If the simulator or runtime experiments produce a negative result, I would document that result in the repository with the same rigor. The contribution would still leave Colibri with reusable traces, a cost model, and evidence preventing the community from repeating an unproductive direction.

Relationship to existing discussions

This connects several existing lines of work without replacing them:

• Full tiered expert streaming #272: full tiered expert streaming;
• Creating a "master expert cache" for new users #282: reusable expert histories;
• [Experiment] expert-transition-history placement policy vs gate-momentum — controlled A/B for hypothesis #1 #708: held-out placement-policy experiments;
• KV cache is stored in fp32 and kvb_all materialises the decompressed K/V (x3.2 vs llama.cpp) #768: KV memory efficiency;
• OLMoE has no hot expert zone (counter-datapoint for pinned hot-store policies) #864: models without a strong global hot zone;
• existing GLM cache, DeepSeek ExpertStore, and Kimi, Inkling, and OLMoE LRU implementations.

The difference is that this proposal asks for one mathematical residency objective and one reusable lifecycle contract, rather than another model-specific cache or prefetch heuristic.

Questions

@JustVugg, would you support Phase 0 as a research contribution before any runtime refactor?

If it demonstrates a robust benefit and the direction is approved, I would be willing to lead the implementation and integration as an upstream contribution, with all code, experimental evidence, and documentation living in this repository under the project's existing governance. No strings attached; I would simply do the work.

More specifically:

1. Does critical-path time avoided per resident byte match the project's intended optimization target?
2. Should DeepSeek V4's ExpertStore lease contract be treated as the starting interface?
3. Would you prefer expert residency and KV-prefix residency in one RFC or split after the simulator stage?
4. Is the incremental adapter approach acceptable for preserving the one-C-file-per-model structure?
5. Which two model families would you want represented in the first simulator traces so the policy cannot overfit to GLM?

If the simulator produces no robust held-out gain, publishing that negative result should be considered the completion of Phase 0 rather than a failed contribution.

[terrizoaguimor]

@JustVugg, before any proposal acceptance or runtime refactor, I completed a real-trace Phase 0 pilot on a DigitalOcean H200 and wanted to add the evidence here. The result is encouraging, but I am keeping the runtime claim deliberately conservative.

What was tested

• GLM-5.2 gs64 int4 container: mastouri/GLM-5.2-colibri-int4-g64-with-int8-mtp
• Model revision: 95bb5f03e3e0ca16b4c711394d0461fa86a3cfcb
• Engine baseline: source at 8d06bd1; Colibri reports v1.5.0
• 24 physical CPU cores, 247 GB RAM, NVIDIA H200 with 143.8 GiB VRAM
• 5 TB local scratch volume mounted at /mnt/scratch
• 429.3 GB model, 142 safetensors shards
• CUDA sm_90 build and coli doctor --deep: passed

I collected 13 real routing traces: 8 training traces and 5 held-out traces covering coding, chat, reasoning, multilingual, and long-context prompts. MTP, cache-aware routing, speculative/pilot loads, and repinning were disabled during collection so the traces describe the ordinary GLM CPU-routing path.

Artifacts were kept outside the repository on the ephemeral server:

/mnt/scratch/traces/train/
/mnt/scratch/traces/eval/
/mnt/scratch/logs/

The offline simulator and its tests are in the local experiment worktree:

c/tools/residency_sim.py
c/tests/test_residency_sim.py
docs/experiments/cnre-offline-simulator.md

Held-out simulator result

The comparison is against uniform LRU under the same expert-residency byte budget. Categories are equally weighted, and the gate requires at least 10% mean felt-wait improvement with no category more than 3% worse.

Expert budget	Candidate	Mean held-out gain	Worst category
120 GB	Uniform frequency admission	+39.91%	+34.84%
120 GB	Dynamic frequency admission	+40.45%	+35.41%
160 GB	Uniform frequency admission	+53.35%	+49.92%
160 GB	Dynamic frequency admission	+54.50%	+50.66%
200 GB	Uniform frequency admission	+65.31%	+62.64%
200 GB	Dynamic frequency admission	+66.29%	+63.07%

The LRU dynamic allocator was effectively neutral in this pilot, around 0.0-0.4% improvement, and did not meet the 10% gate. The frequency candidates passed the nominal category gate at all three budgets.

I also perturbed felt cost independently at representative layers 3, 30, and 60 by 0.5x and 1.5x. The frequency candidates remained above the gate. Across the five held-out PROF=1 logs, runtime instrumentation reported 19.6 seconds of aggregate felt wait over 2,625 loads, approximately 7.47 ms per load. That is an aggregate calibration signal only; current GLM PROF output does not expose per-layer felt wait.

Runtime qualification

The simulator result is not yet a runtime speedup claim. The current GLM engine still admits every demand miss and promotes the bounded tail of each 64-expert block. No runtime frequency-admission implementation was added before posting this update.

The first H200 runtime checks show why this distinction matters:

Configuration	Runtime hit rate	Decode
RAM_GB=120, PIPE=0	57.2%	1.66 tok/s
RAM_GB=200, PIPE=0	73.5%	1.31 tok/s
RAM_GB=120, PIPE=1	57.2%	2.00 tok/s

These are pilot measurements, not an N>=5 ABBA benchmark. They demonstrate that hit rate alone is not the objective: existing I/O/compute overlap changed tok/s more than the cache-size change. An unrecognized COLI_CACHE_ADMISSION=frequency variable had no effect, as expected, because no such runtime policy exists yet.

Interpretation

This is the strongest result so far:

1. The proposed objective is not only a synthetic artifact. Real GLM routing traces show a large and consistent held-out opportunity for admission/frequency policies under constrained expert RAM.
2. The result survives category separation and representative cost perturbations.
3. The result does not justify claiming a runtime gain yet. Felt wait, read overlap, cache churn, and compute contention must be measured after an opt-in implementation.
4. The next implementation should therefore be small and GLM-only: an opt-in demand-admission policy, with no routing change, no new prefetch, no output change, and the existing memory budget held fixed.

If you support continuing Phase 0, I suggest the next reviewable step is that narrowly scoped GLM prototype followed by interleaved ABBA runs reporting tok/s, p50/p99 latency, felt wait, physical bytes read, admissions, rejections, evictions, migration bytes, and exact output identity. If it fails that runtime gate, I will document the negative result rather than promote the policy.

[terrizoaguimor]

Post-merge review of the Phase 0 simulator found three methodological gaps that should be corrected before interpreting the pilot as evidence for runtime admission. I opened the two implementation defects separately so they stay reviewable:

• CNRE simulator calibrates felt wait from loads/token instead of physical misses #924: profile_costs() divides felt-wait seconds by loads/token, a rate, rather than by the physical miss count. The published aggregate us/load calibration from that parser has the wrong denominator.
• CNRE dynamic allocator skips feasible capacity breakpoints #925: the dynamic allocator scans fixed capacity breakpoints and can miss a feasible intermediate capacity that is the first one with positive replay benefit.

The third gap is experimental rather than a parser bug: FrequencyAdmissionPolicy starts evaluation fully seeded from training counts, while the uniform LRU baseline starts empty. That comparison mixes admission behavior with a pretrained resident set. It does not prove the measured gain is false, but it prevents attributing the whole difference to demand admission.

Before any runtime prototype, I suggest adding paired arms with equivalent initial state:

1. cold LRU vs cold frequency admission;
2. LRU and frequency with the same initial residents;
3. separate reporting for seed hits, post-demand admitted hits, startup bytes, and steady-state churn.

The runtime checks already showed that hit rate alone is not the objective, so keeping this qualification strict is useful. I am not proposing runtime code while these offline comparisons remain confounded.