/docs/design/v1/distributed/l2_adapters/overall.md

L2 Store and Prefetch Controller Design

This document describes how the StoreController and PrefetchController interact with L2 adapters, the invariants they maintain, and the assumptions they rely on. It is intended for developers implementing new L2 adapters or modifying the controller logic.

Architecture Overview

                    ┌────────────────────────┐
                    │    StorageManager       │
                    │  submit_prefetch_task   │
                    │  query_prefetch_status  │
                    │  reserve/finish_write   │
                    └────┬──────────┬─────────┘
                         │          │
              ┌──────────┘          └──────────┐
              ▼                                ▼
   ┌────────────────────┐           ┌────────────────────┐
   │  StoreController   │           │ PrefetchController  │
   │  (background thread)│          │  (background thread) │
   │                    │           │                     │
   │  L1 write done     │           │  external submit    │
   │   → store to L2    │           │   → lookup L2       │
   │   → release locks  │           │   → plan + load     │
   └────┬───────────────┘           │   → read-lock L1    │
        │                           └──────┬──────────────┘
        │                                  │
        ▼                                  ▼
   ┌─────────────────────────────────────────────┐
   │           L2AdapterInterface(s)             │
   │   store / lookup_and_lock / load / unlock   │
   │                                             │
   │  Each adapter has 3 distinct event fds:     │
   │   store_efd, lookup_efd, load_efd           │
   └─────────────────────────────────────────────┘

Both controllers run a single background thread each, using select.poll() on eventfds for event-driven I/O. They share the same set of L2 adapter instances (thread-safe by contract) but use different eventfds.

L2 Adapter Interface

L2AdapterInterface (l2_adapters/base.py) provides the non-blocking I/O primitives that both controllers call. All operations follow a submit → poll eventfd → query result pattern.

Event Fds

Each adapter exposes three distinct eventfds:

Method	Used by	Signaled when
get_store_event_fd()	StoreController	A store task completes
get_lookup_and_lock_event_fd()	PrefetchController	A lookup task completes
get_load_event_fd()	PrefetchController	A load task completes

Critical invariant: All event fds across all adapters must be globally unique. The controllers build fd → adapter_index maps; duplicate fds would silently misroute events.

Store Operations

submit_store_task(keys, objects) -> L2TaskId
pop_completed_store_tasks() -> dict[L2TaskId, L2StoreResult]

• Caller provides buffers: The objects list contains MemoryObj references managed by the caller (StoreController holds L1 read locks on them).
• Coarse-grained errors: A store task either fully succeeds or fully fails. The completion dict maps each task id to an L2StoreResult (lmcache.v1.distributed.internal_api.L2StoreResult) that encodes both the success flag and the bytes actually transferred. Use result.is_successful() to check the outcome and result.bytes_transferred() to read the real byte count written to L2 (always 0 on failure). Adapters that fast-path duplicate keys (e.g. skip the write when the key already exists in the backend) should report the real, non-skipped byte count here so the L2 throughput histogram reflects actual work — not submitted-but-skipped bytes.
• Pop semantics: pop_completed_store_tasks() drains all completed tasks. Each task appears exactly once.

Lookup and Lock Operations

submit_lookup_and_lock_task(keys) -> L2TaskId
query_lookup_and_lock_result(task_id) -> Bitmap | None
submit_unlock(keys) -> None

• Locking: lookup_and_lock atomically checks which keys exist and acquires L2-side locks on found keys. This prevents L2 eviction between lookup and load.
• Fine-grained results: Returns a Bitmap where bit i is set if keys[i] was found and locked.
• One-shot query: query_lookup_and_lock_result returns None while pending, then the Bitmap exactly once. Subsequent calls return None.
• Unlock contract: submit_unlock is fire-and-forget. The adapter must guarantee eventual success (retry internally if needed). The caller will never retry.

Load Operations

submit_load_task(keys, objects) -> L2TaskId
query_load_result(task_id) -> Bitmap | None

• Caller provides buffers: The objects list contains pre-allocated L1 write buffers. The adapter writes loaded data directly into these buffers.
• Fine-grained results: Returns a Bitmap where bit i is set if keys[i] was successfully loaded.
• One-shot query: Same semantics as lookup — returns the Bitmap exactly once.

Thread Safety

The adapter must be safe for concurrent calls from the StoreController thread and the PrefetchController thread. In practice, the store operations and lookup/load operations use separate internal state, so this is usually straightforward with per-operation locks or lock-free queues.

Task ID Scope

L2TaskId values are only unique within a single adapter. When tracking tasks across multiple adapters, use the composite key (adapter_index, task_id).

StoreController

Purpose: Asynchronously replicate L1 data to L2 after writes complete.

Source: storage_controllers/store_controller.py

Lifecycle

StorageManager.__init__
  → StoreController(l1_manager, l2_adapters, descriptors, policy)
  → controller.start()        # spawns background thread
  ...
StorageManager.close()
  → controller.stop()         # joins thread, releases locks

Event-Driven Loop

The StoreController's background thread polls on:

1. StoreListener eventfd — fired by L1Manager when finish_write() completes. The listener is an L1ManagerListener registered with L1Manager.
2. Per-adapter store eventfds — fired when L2 store tasks complete.

Data Flow

L1 finish_write()
  │
  ▼ (L1Manager listener callback, inside L1 lock — must be non-blocking)
StoreListener.on_l1_keys_write_finished(keys)
  │  appends keys + signals eventfd
  ▼
_store_loop: poll wakes up
  │
  ▼
_process_new_keys(keys)
  │
  ├─ 1. Group keys by shape (today: (model_name, kv_rank); L1 is a shared
  │     pool so one drain may span models/parallelism configs with different
  │     KV shapes, and each submit_store_task must see uniform (shape, dtype)).
  │
  ├─ 2. For each per-shape group:
  │     StorePolicy.select_store_targets(group_keys, adapters)
  │       → dict[adapter_index, list[ObjectKey]]
  │
  ├─ 3. For each adapter target:
  │     L1Manager.reserve_read(target_keys)  → get MemoryObj + read lock
  │     adapter.submit_store_task(keys, objs)
  │     Track as InFlightStoreTask
  │
  ▼ (later, for each adapter whose store_efd signaled)
_drain_l2_store_completions(signaled_adapters)
  │  adapter.pop_completed_store_tasks() → deposit L2StoreResult
  │  (success flag + bytes_transferred) on each InFlightStoreTask
  │
  ▼
_advance_request(task_key, task)  [state transition]
  │  skip if l2_store_result still None
  │
  ▼
_finalize_store(task_key, task)  [terminal execution]
  │
  ├─ 4. L1Manager.finish_read(read_locked_keys)  → release read locks
  │
  ├─ 5. If success: StorePolicy.select_l1_deletions(keys) → delete from L1
  │     If failure: log warning (best-effort, no retry)
  │
  ▼
Done. Keys remain in L1 unless the policy deletes them.

Lock Invariants

Phase	L1 Lock State	L2 Lock State
Before store	Unlocked	N/A
During store	Read-locked	N/A
After store	Unlocked	N/A

• Read locks during store prevent eviction from removing L1 data while the adapter is reading it.
• Always released: stop() calls _cleanup_in_flight_tasks() which releases all in-flight read locks, even if tasks haven't completed.

StorePolicy

The policy decides two things:

1. select_store_targets(keys, adapters) → dict[int, list[ObjectKey]] Which adapters get which keys. A key can go to multiple adapters. DefaultStorePolicy: all keys → all adapters.

2. select_l1_deletions(keys) → list[ObjectKey] Which keys to evict from L1 after successful L2 store. DefaultStorePolicy: never delete (empty list).

Policies are selected by name via --l2-store-policy (default: "default"). New policies self-register with register_store_policy(name, cls) at import time and are auto-discovered by storage_controllers/__init__.py.

PrefetchController

Purpose: Asynchronously load KV cache data from L2 into L1 ahead of a serving request. Called by StorageManager.submit_prefetch_task() for keys not already in L1.

Source: storage_controllers/prefetch_controller.py

Lifecycle

StorageManager.__init__
  → PrefetchController(l1_manager, l2_adapters, descriptors, policy)
  → controller.start()        # spawns background thread
  ...
StorageManager.close()
  → controller.stop()         # joins thread, releases all locks

External API (Thread-Safe)

# Called from the serving thread
request_id = controller.submit_prefetch_request(keys, layout_desc)

# Polled by the serving thread
result = controller.query_prefetch_result(request_id)  # int | None

• submit_prefetch_request enqueues the request and signals the background thread via an eventfd. Returns immediately.
• query_prefetch_result returns None while in-progress, then the prefix hit count exactly once (pop semantics).

Prefix-Only Loading

Key invariant: Only the contiguous prefix of found keys is loaded.

If L2 has keys {0, 1, 3, 4} but not key 2, only keys {0, 1} are loaded. The gap at index 2 means the vLLM engine cannot use keys 3 and 4 (it needs a contiguous prefix of computed KV cache). Loading them would waste I/O bandwidth and L1 memory.

This is enforced by trim_load_plan_to_prefix() after the policy computes the raw load plan.

Event-Driven Loop

The PrefetchController's background thread polls on:

1. Submission eventfd — signaled by submit_prefetch_request().
2. Per-adapter lookup eventfds — signaled when lookup tasks complete.
3. Per-adapter load eventfds — signaled when load tasks complete.

Request State Machine

Each request goes through two phases:

LOOKUP ──────────────────────────► PLAN_AND_LOAD ──────────► COMPLETED
  │                                    │
  │ submit lookup_and_lock             │ compute load plan
  │ to ALL adapters                    │ reserve L1 write buffers
  │                                    │ submit load tasks
  │ wait for all lookups               │ wait for all loads
  │ to complete                        │ finalize
  ▼                                    ▼

Data Flow

submit_prefetch_request(keys, layout_desc)
  │
  ▼ (cross-thread: submission queue + eventfd signal)
_drain_submission_queue → _pending_queue
  │
  ▼ (if below max_in_flight)
_start_lookup_phase(request_id, keys, layout_desc)
  │
  ├─ Submit lookup_and_lock_task(keys) to EVERY adapter
  │
  ▼ (wait for all adapter lookups to complete)
_advance_request(request, signaled_adapters)  [LOOKUP branch]
  │  _poll_lookup_results(request, signaled_adapters[LOOKUP])
  │  when all_lookups_done():
  │
  ▼
_transition_to_load_phase(request)
  │
  ├─ 1. PrefetchPolicy.select_load_plan(keys, lookup_results, adapters)
  │     → dict[adapter_index, Bitmap]
  │
  ├─ 2. trim_load_plan_to_prefix()
  │     → only keep contiguous prefix keys
  │
  ├─ 3. L1Manager.reserve_write(keys, is_temporary=True, mode="new")
  │     → allocate L1 write buffers
  │
  ├─ 4. Re-trim plan to only successfully reserved keys
  │
  ├─ 5. Phase 1 unlock: unlock L2 keys locked in lookup but NOT in load plan
  │
  ├─ 6. Submit load_task(keys, objs) per adapter
  │
  ▼ (wait for all adapter loads to complete)
_advance_request(request, signaled_adapters)  [PLAN_AND_LOAD branch]
  │  _poll_load_results(request, signaled_adapters[PLAN_AND_LOAD])
  │  when all_loads_done():
  │
  ▼
_finalize_load(request)
  │
  ├─ 7. Phase 2 unlock: unlock all L2 keys in the load plan
  │
  ├─ 8. L1Manager.finish_write_and_reserve_read(loaded_keys)
  │     → atomically: write unlock + read lock
  │
  ├─ 9. L1Manager.finish_write(failed_keys) + delete(failed_keys)
  │     → clean up partial failures
  │
  ├─ 10. Release read locks for loaded keys beyond the prefix
  │      (partial load failures can create gaps)
  │
  ▼
_complete_request(request_id, prefix_hits)
  │  store result, remove from in-flight tracking

Lock Invariants

Phase	L1 Lock State	L2 Lock State
Lookup	None	Found keys are locked
After plan	Write-locked (reserved buffers)	Plan keys locked; others unlocked (phase 1)
During load	Write-locked	Plan keys locked
After load	Read-locked (prefix keys)	All unlocked (phase 2)
After finalize	Read-locked (prefix only)	All unlocked

• L1 write buffers are temporary: is_temporary=True allows the eviction controller to reclaim them if needed, although normally they are short-lived.
• Atomic write→read transition: finish_write_and_reserve_read() ensures no eviction window between write completion and read lock acquisition.
• Post-finalize read locks: The prefix keys remain read-locked in L1 so the serving engine can consume them. The caller (StorageManager) releases these via finish_read_prefetched() after use.

L2 Lock Management

L2 locks prevent adapter-side eviction between lookup and load. They must be released in all cases — success, failure, and shutdown.

Phase 1 unlock (_unlock_unneeded_keys): After the load plan is computed, keys that were locked during lookup but are NOT in the final load plan are unlocked immediately. This happens when:

• The policy assigned a key to a different adapter.
• The key was trimmed by prefix trimming.
• L1 write reservation failed.

Phase 2 unlock (_unlock_all_plan_keys): After load completes (regardless of success or failure), all keys in the load plan are unlocked.

Shutdown cleanup (_cleanup_in_flight_requests): Releases all held L1 and L2 locks for any in-flight requests.

PrefetchPolicy

select_load_plan(keys, lookup_results, adapters) → dict[int, Bitmap]

Receives lookup bitmaps from all adapters and produces a non-overlapping assignment of keys to adapters. Each key appears in at most one adapter's bitmap.

DefaultPrefetchPolicy: For each key, assign it to the first (lowest-indexed) adapter that has it. This is a simple greedy approach.

Policies are selected by name via --l2-prefetch-policy (default: "default"). New policies self-register with register_prefetch_policy(name, cls) at import time and are auto-discovered by storage_controllers/__init__.py.

Max In-Flight Limiting

The controller limits concurrent prefetch requests to max_in_flight (default: 8). Requests beyond this limit are queued in _pending_queue and dequeued as in-flight requests complete.

Note: This is a simple count-based limit. A future improvement would use a dynamic admission controller based on L1 memory usage of in-flight requests.

Integration: StorageManager

StorageManager (storage_manager.py) is the top-level entry point that wires everything together.

Prefetch Flow (from the serving engine's perspective)

# 1. Submit: check L1 first, then delegate remainder to L2
handle = sm.submit_prefetch_task(keys, layout_desc)

# 2. Poll: busy-wait for completion
while True:
    found_count = sm.query_prefetch_status(handle)
    if found_count is not None:
        break

# 3. Read: access the prefetched data (holds read locks)
with sm.read_prefetched_results(keys[:found_count]) as objs:
    # use objs ...
    pass

# 4. Release: drop read locks
sm.finish_read_prefetched(keys[:found_count])

PrefetchHandle

@dataclass(frozen=True)
class PrefetchHandle:
    request_id: int          # -1 if no L2 request needed
    l1_prefix_hit_count: int # leading keys already in L1
    total_requested_keys: int
    submit_time: float       # for latency logging

submit_prefetch_task first checks L1 for a contiguous prefix of hits:

• If all keys hit L1: returns handle with request_id=-1 (no L2 work).
• If some keys miss: submits the remaining keys to PrefetchController.

query_prefetch_status combines L1 hits with L2 results: total_hits = l1_prefix_hit_count + l2_prefix_hits.

Assumptions and Invariants Summary

1. All eventfds are globally unique across all adapters and all operation types. Violating this corrupts the poll-based dispatch.

2. L2 task IDs are per-adapter, not global. Use (adapter_index, task_id) as composite keys.

3. Query results are one-shot. Both query_lookup_and_lock_result() and query_load_result() return a non-None value exactly once per task.

4. submit_unlock must eventually succeed. The controllers will never retry. The adapter must handle retries internally.

5. Prefix-only loading. Only the contiguous prefix of found keys is loaded from L2. Gaps break the prefix.

6. Listener callbacks run inside L1Manager's lock. StoreListener must be non-blocking (append + eventfd signal only). It must never call L1Manager methods (deadlock).

7. L1 write buffers for prefetch are temporary. Allocated with is_temporary=True to allow eviction if needed.

8. Atomic write→read transition. finish_write_and_reserve_read() prevents eviction between completing a prefetch write and acquiring the read lock for the serving engine.

9. Both controllers release all locks on shutdown. stop() always cleans up in-flight tasks, regardless of completion state.

10. L2 adapters are thread-safe. Concurrent calls from the StoreController thread and PrefetchController thread are expected.

Implementing a New L2 Adapter

Pure-Python Adapters

Implement L2AdapterInterface directly. See mock_l2_adapter.py for a reference implementation for an in-memory adapter, or raw_block.md for a durable local-device adapter. No existing files need to be modified. Create a new module (e.g., my_l2_adapter.py) in the l2_adapters/ package and self-register at module level:

# At the bottom of your module:
register_l2_adapter_type("my_type", MyL2AdapterConfig)
register_l2_adapter_factory("my_type", _create_my_l2_adapter)

The __init__.py uses pkgutil.iter_modules() to discover all *_l2_adapter.py modules automatically, but imports them lazily — a module (and its third-party dependencies) is only loaded when the corresponding adapter type is actually requested at runtime.

Persist / Recover (Optional)

Adapters that support persisting cached data across restarts use PersistConfig (parsed from the JSON key "persist_enabled", defaulting to True). Lookup always checks secondary storage on miss. There is no dedicated interface method — adapters integrate persist into their existing close() (to keep data on disk) path. See nixl_store_dynamic_l2_adapter.py for a reference implementation and nixl_store.md for design details.

Native (C++/Rust) Storage Backends

For high-performance backends written in C++ or Rust, use the shared native connector framework. A single C++ connector implementation works in both non-MP mode (via ConnectorClientBase) and MP mode (via NativeConnectorL2Adapter).

Full guide: csrc/storage_backends/README.md

The NativeConnectorL2Adapter (native_connector_l2_adapter.py) bridges any pybind-wrapped IStorageConnector to the L2AdapterInterface:

• Creates 3 Python eventfds from the connector's single eventfd
• Runs a background demux thread that routes completions by operation type
• Handles ObjectKey serialization and MemoryObj buffer extraction
• Implements client-side locking (refcount dict) for remote backends

Reference implementation: The Redis (RESP) connector in csrc/storage_backends/redis/ demonstrates all 5 steps of the integration guide.

Implementing a New Store or Prefetch Policy

Both store and prefetch policies use a name-based registry with automatic module discovery. To add a new policy, create a single file in storage_controllers/ — no changes to any existing file are needed.

Store Policy

1. Create a new file (e.g., storage_controllers/store_policy_tiered.py).
2. Subclass StorePolicy and implement select_store_targets() and select_l1_deletions().
3. Call register_store_policy("tiered", TieredStorePolicy) at module level.

from lmcache.v1.distributed.storage_controllers.store_policy import (
    StorePolicy,
    AdapterDescriptor,
    register_store_policy,
)
from lmcache.v1.distributed.api import ObjectKey


class TieredStorePolicy(StorePolicy):
    def select_store_targets(self, keys, adapters):
        # custom logic ...
        ...

    def select_l1_deletions(self, keys):
        return []


register_store_policy("tiered", TieredStorePolicy)

The policy is now available via --l2-store-policy tiered.

Prefetch Policy

Same pattern: subclass PrefetchPolicy, implement select_load_plan(), and call register_prefetch_policy("name", cls).

How Discovery Works

storage_controllers/__init__.py uses pkgutil.iter_modules() to import every module in the package at import time. When your module is imported, the register_*_policy() call at module level adds it to the registry. The --l2-store-policy and --l2-prefetch-policy CLI arguments use the registry to populate their choices list.