A Token’s Journey Through Colibri

In JustVugg/colibri, a routed expert becomes useful only when a logical decision—“this token needs these experts”—is turned into a safe physical sequence of lookups, reads, computations, and residual updates. This explanation follows that sequence inside one batched MoE call. The goal is to make Colibri’s placement thesis operational: readers should be able to tell what happens on a cache hit, what happens on a disk miss, why a batch union saves work, and where asynchronous execution can hide time without changing the intended route.

The route becomes a batch

A Colibri MoE call begins with hidden-state rows. A row represents one position being processed; it might be one token during decoding or many token positions during prompt processing. The router examines each row and assigns a small set of expert IDs with corresponding gate weights. Those IDs are the logical work request for the layer. At this point, the engine has not yet decided whether an expert will come from RAM, storage, or an accelerator. Routing comes first because placement cannot be chosen intelligently until the request is known.

For the reference GLM engine, this phase can be understood as two related operations. First, it ranks experts for each row and keeps the configured routed set unless an explicit experimental control asks for a reduction. Second, it records the selected IDs and weights so the later computation can reproduce the routed mixture. The weights are not merely a hint about which experts seem interesting. They determine how each expert’s output contributes to that row, so the route and its gate values have to travel together.

A small example shows why this information should be represented for the whole batch. Suppose three rows request these experts:

  • row A: experts 2, 7, and 11;
  • row B: experts 7, 11, and 19;
  • row C: experts 2, 19, and 23.

There are nine routing choices but only five distinct experts. If the engine loads separately for every row, it can read expert 7 twice, expert 11 twice, expert 2 twice, and expert 19 twice. That repeats the same weight movement even though the expert’s matrix is identical for all rows in this call. Colibri instead forms a batch union: it keeps one copy of each distinct expert and remembers which rows use it and with what gate weight.

The union is more than a small optimization. It separates two kinds of multiplicity that would otherwise be confused. A selected expert can serve many rows, so its matrices should be loaded once. Each row still needs its own activation and its own weighted contribution, so the computation must preserve row identity. The union removes duplicate I/O without merging the mathematical contributions of different positions.

The same idea becomes more valuable during prompt processing. A prompt contains many rows, and several rows may converge on the same experts. Colibri can ask for a unique expert set, load each member once, compute the expert’s feed-forward network for the relevant rows, and scatter the results back to their original positions. During single-token decoding, the union may be smaller, but the principle remains the same: deduplicate physical weight movement while retaining logical routing.

The engine also collects routing telemetry around this phase. It counts expert selections for the persistent usage history and can write a trace containing the call, row, layer, expert IDs, and applied gates. That information serves two different purposes. The counters help the learning cache decide which experts deserve residency over many requests. The trace can support later prediction experiments. Neither should be mistaken for a second router: the current call still uses the selected IDs, and ordinary telemetry does not replace the model’s routing function.

There are explicit boundaries here. A top-p or top-k reduction can lower the number of experts and therefore the amount of data the engine must move, but it may change the routed computation. A cache-aware routing experiment can prefer resident experts within a ranking window, but that is different from ordinary placement because it can substitute expert IDs. Colibri’s default path keeps these choices separate from the normal route. The engine can measure an experimental policy, but its output claim must identify whether the policy merely moved selected weights or selected a different set.

After the route and union are ready, the logical request is still incomplete. The list says which experts are needed, not where their bytes are. That is the point at which Colibri’s memory hierarchy becomes an execution decision.

sequenceDiagram participant R as Router participant U as Batch union participant C as Pin/LRU cache participant D as Disk or VRAM/RAM tier participant X as Expert compute participant T as Telemetry R->>R: Select expert IDs and gate weights per row R->>U: Submit routed rows U->>U: Deduplicate expert IDs and retain row references U->>C: Resolve each unique expert alt Pinned or LRU hit C-->>X: Return resident expert view else Disk miss C->>D: Load expert slab into a working slot D-->>C: Return weights and scales C-->>X: Publish loaded expert view end X->>X: Compute expert FFN for its rows X-->>R: Scatter gate-weighted results R->>T: Record hits, misses, and routes

A miss is a controlled state transition

For each unique expert, Colibri checks the pinned hot store first and then the per-layer LRU cache. A pinned hit is an expert loaded into the long-lived hot set; it is intended to remain available until a placement update changes the set. An LRU hit is resident because recent or repeated activity gave it a working slot, but it can be displaced when a colder expert needs that capacity. A disk miss is neither a missing model expert nor a different model choice. It means the selected expert’s bytes are currently outside the fast working set.

That distinction is important when reading a hit-rate line. Two runs can have the same total hit percentage but different behavior if one gets most of its hits from stable pins and the other repeatedly churns the LRU. A pin avoids repeated loading because the policy has committed space to the expert. An LRU hit reflects observed locality that remains reversible. A miss exposes the storage and loading path, and its cost depends on the size of the expert, the drive, the cache state, and whether another computation can run while the read is in flight.

A miss is placed into a reusable working slot. The GLM path stores the expert’s gate, up, and down matrices together when the on-disk layout allows it, along with their scale data. The slot becomes a set of tensor views over the loaded slab rather than a separate permanent allocation for every expert. Once the read completes and the format has been resolved, the expert can be passed to the computation phase. When the slot is later reused, the old expert must no longer be visible through the cache index.

This is where a lease matters. The expert-store interface requires a successful lookup to be paired with exactly one release. A view is valid only while its lease is active. The lease prevents the cache from treating a borrowed slot as free while a compute or device operation is still reading it. A prefetch request, by contrast, is advisory and holds no lease; it must not evict a slot that still has an active lease. These rules make ownership explicit at the boundary between residency management and computation.

The need for that boundary grows when accelerators are involved. A GPU may still be reading a resident expert after the CPU has finished issuing the operation. Colibri tracks in-flight use so a repinning or replacement operation cannot free or overwrite the slab too early. If a slot is reused for another expert, stale device copies are invalidated before the new expert is allowed to compute. The purpose is not abstract concurrency hygiene. Without it, a fast run could combine part of one expert with part of another and return plausible but incorrect numbers.

The ordinary blocking path loads all current misses before the compute thread proceeds. That is simple and gives a clear wait boundary, but it leaves the processor idle during storage service. With PIPE enabled, Colibri dispatches missing loads to a bounded asynchronous pool. Resident experts can then enter computation while the missing slabs are being read. The overlap is real only up to the point where the computation needs the missing expert. The layer still has to wait before it can finish the routed result, because a selected expert cannot be replaced by an empty buffer just to preserve the illusion of progress.

This is why “asynchronous” does not mean “non-blocking inference.” It changes when the wait occurs and how much of it is hidden. If resident expert work takes longer than the read, the load may complete before the compute reaches it and the miss becomes almost invisible. If the disk is slow or the batch contains mostly misses, the final join still exposes the storage delay. Colibri’s profiling separates disk service from foreground-visible wait for precisely this reason: the device may be busy for a long time even when some of that service overlaps with useful computation.

The engine can also issue readahead for the next block of unique experts while the current block computes. That prefetch is a prediction about imminent demand, not a replacement for the current route. It can help when the next block’s reads have time to warm the relevant cache or page data. It can hurt when the drive is already saturated, because speculative traffic competes with the demand read that must complete the current layer. The safe interpretation is therefore “a possible earlier start for a known future request,” not “free bandwidth.”

The same causal structure applies to a resident GPU tier, with a different execution site. If the requested expert is already registered on the accelerator, Colibri can avoid a RAM slot and disk load for that expert, then combine its result with CPU-served experts. On a unified-memory system, resident expert work can be submitted while the CPU handles missed experts. On a discrete GPU system, moving an expert from storage for every token would introduce another transfer bottleneck, so the useful GPU tier is the set that can stay resident and be reused. The tier changes where an already-available expert runs; it does not make a cold expert instant.

Once the needed views are available, the computation follows the logical route. For each selected expert, Colibri applies the expert’s gate and up projections, the activation function, and the down projection. The resulting expert output is multiplied by the route’s gate weight and accumulated into the output row for every position that selected that expert. Because the batch union kept row references, one expert computation can serve many positions without losing their separate weights. When all unique experts have contributed, the routed result is ready to rejoin the layer’s residual stream.

Telemetry closes the loop. The engine increments request, hit, and miss counters; records which tier served a hit where applicable; and updates the usage and recency information used by later placement decisions. A successful computation therefore changes more than the current hidden state. It also changes the evidence the cache will use the next time the model runs. That feedback is useful only if the counters represent the actual route and the slot index still names the bytes it claims to hold.

The failure paths explain why the implementation treats these stages as one contract. If a required tensor is missing, a read ends early, or a loaded payload does not match its expected format, the main demand path does not quietly continue with partial weights. A speculative future-layer load can be abandoned without killing the server because it was only a prediction; the current routed computation can later perform a normal demand load. The difference is ownership of the result: a demand miss is required to complete the selected route, while a speculative load is allowed to fail and leave no semantic trace.

There is also a clean termination condition for the per-layer journey. Every selected expert has either been resolved from a valid resident view or loaded successfully into a valid working slot; every row has received the weighted contributions from its selected experts; the views have been released only after their consumers are finished; and the counters describe what happened. If storage capacity or bandwidth is insufficient, the process may take longer. The route itself should not silently become a different route merely because a cache lookup failed.

This is the practical meaning of Colibri’s “placement changes speed” principle. The router owns the logical request. The cache owns residency and replacement. The storage layer owns the byte-range read. The expert kernel owns the numerical transformation. Telemetry observes the path and informs future placement. Those responsibilities meet at explicit boundaries rather than being collapsed into one operation that could hide whether a problem came from routing, I/O, memory, or compute.

The next layer of the design is the weight store itself: how Colibri finds the three matrices in a many-shard checkpoint, verifies that their bytes and shapes make sense, and decodes them into the quantized forms that the expert computation expects. Without that substrate, the token path would have a logical route but no safe physical source for its weights.

References