[Feature]: Performance Optimization of the Multi-Rank Shared-Cache Load Path

1. Background and Terminology

Model sharding and cache. During large-model inference, the KVCache data produced is large; UCM divides it into equal-sized segments called shards. A "block" is identified by a 16-byte ID, and each segment inside a block has an index, so locating "a given segment of a given block" takes two quantities: (block, intra-block shard index) — this is the cache key. The system maintains a host-memory cache CacheBuffer that holds recently used shards; a hit avoids going back to the source. The cache consists of a fixed number of "slots"; when full and a new shard must be loaded, some policy picks an old slot to free — this is eviction.

Multi-rank sharing. On a multi-GPU node doing tensor parallelism (TP), each card is a rank; the number of cards in the node is localRankSize. Key design: all ranks in the node share the same CacheBuffer — both its control plane (metadata, locks, state, eviction cursor) and data plane (actual bytes) are shared across ranks. The control plane is published via shared memory (each rank maps the same physical page and sees the same locks and metadata); the data plane is built from huge-page-backed segments whose file descriptors are passed to other ranks over a Unix domain socket, so "the same bytes are visible to all ranks and already registered with the device." Each rank runs its own LoadQueue: internally it has a dispatcher (takes tasks from the pending queue, locates/allocates slots in the cache, and goes back to the backend to load when needed) and a copier (once data is ready, DMAs it to device memory).

Owner and waiters. When a shard is requested for the first time and misses in the cache, the first rank to grab that key's lock and allocate a slot becomes the owner, responsible for issuing the backend Load; other ranks requesting the same shard find someone is already filling it (refcount > 0) and become non-owner waiters, not duplicating the load.

Load tasks. The system doesn't load one segment at a time; it batches segments that need loading into one task. A task contains shards of multiple different blocks, each block contributing one shard (each entry looks like "block B, the s-th segment of that block, target device address"). This amortizes task-submission/queue-advancement overhead.

---

2. Problems to Solve

2.1 Lock contention and serialized loading

In the old implementation each rank iterated shards in natural order 0,1,2,…. Picture 4 ranks handling a 12-shard task: at step 0 all ranks request shard 0, at step 1 all request shard 1… at every step, all ranks pile onto the same shard — i.e., the same bucket/node lock protecting it. Locks are exclusive, so dispatch is serialized. Worse, only the one rank that wins owner each step goes back to the backend; the other N−1 ranks idle. The result: the slow backend has only 1 concurrent load; the larger the TP, the more waiting, and throughput drops rather than rises.

2.2 Pure ring eviction hurts locality

The old eviction was a pure FIFO "cursor +1 mod", with no hot/cold distinction. When batch-loading many block-shards, a just-loaded shard about to be read by the copier can be picked as a victim in the next miss's eviction sweep — a wasted load that throws away the hard-won locality.

2.3 Non-owner waiting via polling

In the older implementation, the ready state was just a plain bool in Meta with no wait mechanism, accessed under the node lock. A non-owner rank wanting to know "is this slot done" could only poll: lock — read bool — unlock — sleep a bit — repeat. Three harms: each waiter burns CPU continuously (more waiters as TP grows); after readiness, the waiter only discovers it at the next poll cycle (high wake latency); and a plain bool has no synchronization semantics and isn't visible across processes in shared memory, so cross-process load state was unreliable to begin with.

2.4 Wasted dispatcher idle window

After the dispatcher pushes all shards of the current task to the copier, it waits for the next task to arrive. This idle window could prewarm the "next slot" to cut later arbitration latency, but the old implementation didn't.

---

3. Solutions

3.1 RearrangeIndex: rotate the access sequence #1204

Each rank starts at a different offset, taken as deviceId % localRankSize. The reordered sequence strides by localRankSize from each rank's offset, taking its own stride first, then rotating. Two compounding benefits: (1) different shards land on different bucket/node locks, so Get() no longer serializes in lockstep; (2) all ranks simultaneously become owners of their respective shards, issuing N concurrent loads to the slow backend — ownership is balanced across ranks (≈1/N each), backend bandwidth is saturated, and copier copy volume is balanced too. These are two faces of the same effect: rotation both removes Get()-internal lock contention and turns loading from serial to concurrent.

Scenario A — lockstep contention vs. rotation (4 ranks × 12 shards). Old natural order: at step 0 all 4 ranks request shard 0 (lands in bucket 37), queuing on "bucket 37's lock" — rank 2 wins first, becomes owner, issues a backend Load; the other three register as waiters and go poll; step 0 has only 1 concurrent backend load. Step 1, all pile on shard 1 again… 12 steps with only 1 concurrent backend load throughout. After rotation: step 0 rank0→shard0 (bucket 37), rank1→shard1 (bucket 12), rank2→shard2 (bucket 88), rank3→shard3 (bucket 5) — four different locks, no blocking; all four ranks become owners simultaneously and issue 4 concurrent backend loads. Step 1 rank0→shard4, rank1→shard5, rank2→shard6, rank3→shard7, still 4-way concurrent. In the first 3 steps the four ranks hit mutually disjoint shards each step; after sweeping their own stride they rotate to the next, still keeping "different shards per step." A 12-shard task goes from "12 serial steps" to "3 steps, 4-way concurrent each" — throughput scales linearly with TP.

3.2 CacheBuffer split into control plane + data plane

The formerly unified storage implementation is split into BufferCtrl (control plane) and BufferData (data plane): the control plane holds metadata/locks/state/bucket-chains/eviction-cursor/waiting, compact and hot; the data plane holds bytes, large and huge-page-aligned. Each has its own layout and memory strategy. The deviceId<0 "query-only, no-copy" role (Scheduler) builds only the control plane, no data plane; DataAt returns null and only answers "exists." Relative to the old implementation, this regional refactor adds several performance points:

• Single state machine: merge the two separate ready/failed bools into one State (Idle/Ready/Failed), reducing status determination from two atomic reads to one, and explicitly introducing the Idle state to make "can-takeover" semantics clear.
• Wait interface returns state directly: previously the caller had to re-check ready/failed (one read each) after waiting; the new interface returns the state, so one branch suffices — this also naturally evolves the LoadQueue wait loop from "poll repeatedly" to "branch on the returned state."
• Slot-addressing fast path: locating "slot i in which segment, at what offset" used division + modulo throughout; the new implementation checks whether nodes-per-segment is a power of two, and if so uses shift + mask instead of division (equivalent but much faster), otherwise falling back to div/mod. Segment sizes are usually engineered to powers of two, so the fast path is the norm.
• Non-owner waiting from polling to futex block/wake: ready/failed become cross-process-visible atomics, with futex — non-owners no longer poll but block (zero CPU); on completion the owner writes the state, then increments a version, then wakes. The version is the key to no-lost-wakeup: the waiter remembers the current version before blocking; any state change bumps the version before waking, so no wake can be missed between "check state" and "block." Benefits: waiters don't burn CPU, readiness wakes instantly, and cross-process state visibility is filled in.

Scenario B — owner/non-owner and futex wake. Shard X is won as owner by rank 0 at step 0. Old polling: rank 2 later requests X, gets the bucket lock, sees ref>0, registers as waiter, releases the lock, and enters the "lock node — read ready bool — unlock — sleep 10ms" poll; rank 0 fetches X and writes ready=true; rank 2 reads it at most 10ms later and exits — those 10ms are either busy-spinning or waking late, and 4 waiters mean 4 wastes. New futex: rank 2 registers as waiter, then calls "wait-for-change" — reads the current version g0 and state, and if not ready blocks (yields the CPU); rank 0 finishes, writes state=Ready, bumps the version to g1, and wakes; rank 2 wakes instantly, reads Ready, and exits. If rank 0 happened to finish between rank 2's "read g0" and "block," since finishing always bumps the version before waking, rank 2's "wait on the same g0" returns immediately — it can't miss this completion.

3.3 CLOCK second-chance eviction + block-level Touch #1274

Replace the pure ring with CLOCK: each slot gets an "accessed bit"; the cursor scans slot by slot — bit=1 is cleared and skipped (a "second chance"), bit=0 is chosen as the victim; if all bits are set, it degenerates back to the raw cursor to guarantee termination. Bits are set on both "hit" and "new-load": the former keeps hot data alive, the latter gives freshly-loaded data a reprieve window so it isn't swept before the copier reads it — critical for batch-loading many block-shards, preventing the "load-then-immediately-evict" livelock. Block-level Touch refreshes a batch of blocks' hotness in one shot — it walks the buckets they land on, setting the accessed bit on every resident slot (all resident intra-block shards of a block get marked), and per the cascading contract forwards the same batch to the backend, so that if this layer later evicts a block, the lower layer still holds a warm copy and reloading is served by the lower layer — no thrashing.

Scenario C — CLOCK preserves a hit block. Cache of 4 slots, loaded A/B/C/D in turn; on load each slot's accessed bit is set to 1, cursor at slot 0. A miss (E) arrives: slot0 bit=1→clear, skip; slot1 bit=1→clear; slot2 bit=1→clear; slot3 bit=1→clear; one lap clears all to 0, degenerates to the raw cursor, picks slot0 as victim, E goes into slot0. Now slot0=E (bit=1, just loaded), slots1/2/3=B/C/D (bit=0, just swept), cursor at slot1. Someone Gets B (slot1) — the "hit" sets slot1's bit back to 1. Another miss (F): cursor from slot1, slot1 bit=1→clear, skip (B survives); slot2 bit=0→chosen as victim, F goes into slot2. The just-touched B survives, the never-touched C is evicted — CLOCK approximates LRU, avoiding "evicted right after use."

Scenario D — Touch preserves all shards of a block. Cache of 6 slots; block T's 3 intra-block shards T.s0/s1/s2 are in slots 0/1/2, the other 3 slots hold U/V/W. First, one miss sweeps all the "new-load" bits clear; now T's three shards are in slots 0/1/2, all bits=0. The caller anticipates T will be densely accessed soon and calls Touch on T: clustered by bucket, walk the bucket chains, find all of T's resident shards (slots 0/1/2), set all three bits to 1. A miss (X) arrives: cursor from slot0, slot0 bit=1→clear, skip (T.s0 survives); slot1 bit=1→clear, skip (T.s1 survives); slot2 bit=1→clear, skip (T.s2 survives); slot3 bit=0→chosen as victim, X goes into slot3. All three shards of block T survive, and the unrelated U is evicted — block-level Touch renews a whole block as a unit.

3.4 Idle-window preallocation #1274

Add Prealloc: after pushing all shards of the current task to the copier, in the idle window before the next task the dispatcher runs an extra loop prewarming the next intra-block shard slot for each block — assuming intra-block shards are streamed in order, the next task likely wants each block's next slot. Prewarm: allocate an idle slot for (block, next-index), hang it on the bucket chain, drop the refcount to 0 (slot stays, state idle); a later Get takes the "hit-or-takeover" fast path, skipping the "evict—unlink—rehang" retry; if the regular pool is full, the reserved pool is used, not crowding normal capacity. A next-index that doesn't exist yet just leaves an idle slot, harmless.

Scenario E — prewarming the next slot. A task has B0/B1/B2 each at segment s; the dispatcher pushes B0.s/B1.s/B2.s to the copier in RearrangeIndex order, then in the idle window prewarms (B0,s+1), (B1,s+1), (B2,s+1) — each an idle slot on the bucket chain, ref=0; the regular pool full → use the reserved pool. The next task requests (B0,s+1): Get finds it already on the chain (the prewarmed idle slot) with ref==0 + idle, takes over as owner directly, skipping "find a victim + cross-bucket trylock-retry." If the prewarmed slot was already taken by someone else (concurrency), it sees ref>0 and registers as a waiter via the fast path. For B3, which isn't in this task, nothing happens — prewarming only acts on the current task's blocks; the scope is bounded.

---

4. Benefits Summary

• Throughput scales with TP: rotation removes lockstep contention + balances owners → backend concurrency from 1 to N.
• Cache hit rate & locality: CLOCK ≈ LRU, keeps just-loaded/just-hit shards; Touch lets hot spots actively renew; cascading ensures eviction doesn't cause reload thrash.
• CPU & latency: non-owner waiting from polling to block/wake, saving CPU and waking instantly.
• Dispatch latency: the idle window prewarms the next slot, shortening later Get arbitration.