# NSA Prefill CP HiCache Async State Synchronization Notes

Date: 2026-05-27

Branch context: `cp-hicache-refactor` around `d14c02b0dc`.

This note records the current understanding of CP shared-KV + HiCache
backup/load/evict behavior and the synchronization model needed to make it
both correct and performant.

## Problem Statement

With CP shared KV enabled, the radix cache and scheduler operate in a full
logical KV address space, while every CP rank owns only a physical shard of the
device and host KV pools. HiCache backup, load, and eviction therefore have two
different state domains:

- **Logical state** must be identical on every CP rank.
  - Radix tree topology.
  - Node residency state.
  - `TreeNode.value` logical loc tensor.
  - `host_len`, `cp_hicache`, `page_owners`.
  - Lock/ref counters that affect eviction eligibility.
  - Request prefix indices and scheduler admission decisions.
- **Physical state** is intentionally rank-local.
  - Physical device pages owned by the local CP rank.
  - Physical host slots allocated by the local CP rank.
  - Zero-owned ranks for a given logical node.
  - Draft/MTP host slots for locally owned pages.

The correct model is therefore:

> Logical operations must be submitted and committed in the same order on all
> CP ranks; physical transfers may be sparse and rank-local.

The current code partially implements this, but the async lifecycle is still
spread across `HiRadixCache`, `HiCacheController`, allocator state, and CUDA
events. This makes performance tuning risky because reducing synchronization
can easily introduce cross-rank state divergence.

## Target Contract

The intended design for the current CP HiCache stage is:

1. **Radix tree is synchronous.**
   - All CP ranks apply radix match/insert/split/evict and residency metadata
     changes in the same order.
   - Background transfer completion must not independently mutate radix
     topology or make partially transferred nodes visible.

2. **Physical transfer is asynchronous.**
   - D2H backup and H2D load may run on transfer streams and overlap with
     scheduler progress or layer compute.
   - Async completion only satisfies an already-submitted logical operation.
     It does not decide which logical node exists, which node is visible, or
     which victim is chosen.

3. **Draft/MTP KV is strongly consistent with target KV.**
   - Draft KV is a shadow payload of the target radix node.
   - Target and draft residency transition together: host valid, device valid,
     eviction, load, backup, split, and rollback are all atomic at the logical
     node level.
   - If target and draft disagree, the node must be treated as invalid for
     speculative decoding.
   - `--hicache-size` is the per-rank host budget for the whole logical cache
     object, not a budget to be independently duplicated by target and draft.
     When draft/MTP HiCache is attached, target and draft host pools must be
     sized from this shared budget. Draft may have a larger token capacity when
     its per-token footprint is smaller, but it must follow target residency
     and allocation decisions.

4. **Cache-hit load is layer-wise.**
   - On host hit, device allocation and radix state are resolved
     synchronously, then each layer's target/draft payload is loaded from host
     to device asynchronously.
   - The model waits per layer through the existing layer-ready mechanism.
   - The request may not observe a layer before both target and draft payloads
     for that layer are ready.

5. **Backup should become layer-wise.**
   - After each layer stores its target/draft KV for the computed suffix, that
     layer can be backed up to host asynchronously.
   - The radix node becomes host-visible only after all required layer backups
     have completed.
   - This keeps the logical behavior close to `write_through`, but reduces the
     tail cost by overlapping D2H with remaining layer compute.

In short:

```text
sync radix/control plane
async transfer/data plane
target KV and draft KV are one logical cache object
per-layer host load on hit
per-layer host backup after layer KV store
```

## Current Code Facts

### CP owner mapping

`CpSharedKVLayout` defines page ownership by logical page id:

- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`
  - `owner_for_logical_pages()`
  - `owned_by_this_rank()`
  - `logical_locs_to_physical()`

Logical page 0 is padding. Real pages are distributed by:

```text
owner = (logical_page - 1) % cp_size
physical_page = (logical_page - 1) // cp_size + 1
```

This is the basis for all CP HiCache physical filtering.

### Backup path

Main paths:

- `HiRadixCache.write_backup()`
  - `python/sglang/srt/mem_cache/hiradix_cache.py`
- `HiCacheController._write_cp()`
  - `python/sglang/srt/managers/cache_controller.py`
- `HiCacheController.start_writing()`
  - `python/sglang/srt/managers/cache_controller.py`

Current behavior:

1. `TreeNode.value` is full logical locs.
2. `_write_cp()` filters local owned positions.
3. Local logical locs are translated to physical device locs.
4. Host pool allocation is only for local physical slots.
5. Target and draft host slots are reserved together when draft HiCache is
   registered.
6. D2H is submitted on `write_stream`.
7. Zero-owned ranks enqueue a completed no-op write ack.
8. `writing_check()` uses `all_reduce(MIN)` over locally completed ack count so
   all ranks pop the same logical write ack prefix.

This gives asynchronous data movement, but logical backup state is updated in
`write_backup()` before a full CP-group commit protocol exists.

Another important asymmetry: backup currently calls
`backup_from_device_all_layer()`, while load uses `load_to_device_per_layer()`.
So load has layer-wise overlap with compute, but backup is submitted as an
all-layer D2H operation after the node has been inserted. The current backup
path can overlap with later scheduler work, but it does not overlap each
already-finished layer's D2H with the remaining layers of the same prefill
forward.

### Load path

Main paths:

- `HiRadixCache.load_back()`
  - `python/sglang/srt/mem_cache/hiradix_cache.py`
- `HiCacheController.load_cp()`
  - `python/sglang/srt/managers/cache_controller.py`
- `HiCacheController.start_loading()`
  - `python/sglang/srt/managers/cache_controller.py`
- `LayerDoneCounter`
  - `python/sglang/srt/managers/cache_controller.py`
- KV pool layer access waits
  - `python/sglang/srt/mem_cache/memory_pool.py`

Current behavior:

1. Load-back walks evicted host-backed radix nodes.
2. CP mode reconstructs the write-time owner pattern from
   `node.cp_hicache.page_owners`.
3. `alloc_pages_with_owners()` allocates fresh logical pages whose owner lanes
   match the backed-up metadata.
4. Only local owned positions are transferred from host to device.
5. H2D load is submitted per layer on `load_stream`.
6. `LayerDoneCounter` records a producer id.
7. Forward receives `hicache_consumer_index`.
8. Attention reads call `wait_until(layer_id)`, allowing layer-wise load/compute
   overlap instead of blocking on the full load.

This is the strongest current async mechanism: load can overlap with prefill
compute at layer granularity.

### Device eviction

Main path:

- `HiRadixCache.evict()`
  - `python/sglang/srt/mem_cache/hiradix_cache.py`

Current behavior:

- If a node is already host-backed, device eviction frees device slots
  immediately and keeps host metadata.
- If `write_policy == "write_back"` and a node is not host-backed,
  `write_backup(..., write_back=True)` is called and then
  `writing_check(write_back=True)` synchronizes all write completion before
  freeing device memory.

This means write-back eviction is not truly asynchronous in the hot allocation
path. It can enqueue D2H asynchronously, but then blocks before device memory is
released.

### Host eviction

Main path:

- `HiRadixCache._evict_host_for_physical_slots()`
  - `python/sglang/srt/mem_cache/hiradix_cache.py`

Current behavior:

- Victims are selected from logical host leaves.
- Each rank frees its own physical host slots through `evict_cp_host()`.
- With `synchronize_across_ranks=True`, a loop calls `all_reduce(MIN)` on a
  local done flag until all ranks have freed enough local physical host slots.

This preserves CP correctness under uneven per-rank host pressure, but the
current loop can produce frequent CPU all-reduces when host memory is near
capacity.

## Key Invariants to Preserve

1. **Radix logical state is replicated.**
   - All ranks must insert, split, evict, load, and mark backup state for the
     same logical nodes in the same order.

2. **Physical ownership is rank-local.**
   - A rank transfers only pages it owns.
   - Zero-owned ranks still participate in logical op ordering with no-op
     acks.

3. **Host metadata must preserve write-time owner pattern.**
   - `page_owners` is required because a later load allocates fresh logical
     pages. The new allocation must reproduce the write-time owner pattern or
     `owned_positions` will refer to wrong physical shards.

4. **Target and draft KV must be atomic.**
   - For EAGLE/MTP, target KV and draft KV must be backed up, loaded, and
     evicted under the same logical operation. Partial target-only or
     draft-only success can corrupt accept rate after cache hit.

5. **Async transfer source/destination lifetime must be protected.**
   - A device node under D2H write must not be evicted until the write ack is
     committed.
   - A device node under H2D load must not be used by a layer until that layer's
     load event is complete.

6. **Failure must be globally resolved.**
   - If any rank cannot reserve host slots, reserve draft slots, allocate owner
     lanes, or submit a needed transfer, all ranks must either commit or abort
     the logical operation together.

## Current Design Boundary

The first stabilization pass should intentionally keep the control plane small.
Do not include TBO, decode-side behavior, L3/storage, or cross-instance cache
sharing in this stage. The scope is:

- prefill-side CP shared KV,
- L1 device KV and L2 host HiCache,
- target KV plus EAGLE/MTP draft KV,
- synchronous radix tree mutation,
- asynchronous low-level D2H/H2D data movement.

The control-plane rule is:

> Radix tree operations are synchronous; only physical KV movement is async.

The synchronization rule is:

> Prefer deterministic-by-construction state transitions over explicit
> all-rank confirmation.

This means the scheduler/radix layer remains the sole owner of logical cache
state. CUDA events and background transfer callbacks may report progress, but
they must not independently mutate radix topology, node residency, evictable
sets, or request-visible cache-hit state.

Synchronous control-plane operations include:

- radix match, insert, split, and eviction metadata updates,
- host/device residency state transitions,
- target/draft metadata commit or rollback,
- logical lock/ref updates,
- evictable leaf set updates,
- device and host slot reservation decisions,
- global commit/abort decisions after slow-path failure.

Asynchronous data-plane operations include:

- target KV D2H backup,
- draft KV D2H backup,
- target KV H2D load,
- draft KV H2D load,
- layer-wise load readiness events.

This is deliberately conservative. It preserves CP rank determinism while still
allowing backup and load transfers to overlap with scheduler progress or model
compute.

## Synchronization Philosophy

A healthy CP HiCache design should not need all ranks to explicitly vote on
every normal cache transition. The common path should be correct because every
rank runs the same deterministic logical state machine:

```text
same request order
same radix mutations
same logical page ids
same owner pattern
same victim selection
same target/draft atomic op boundaries
```

Given those inputs, each rank can independently derive:

- which logical node is affected,
- which physical pages it owns,
- how many local host/device slots are required,
- which local slots can be freed,
- whether it is a zero-owned participant,
- which local async event corresponds to the logical op.

All-reduce should therefore be a **slow-path safety mechanism**, not the normal
source of truth. It is appropriate only when a condition is genuinely
rank-local and cannot be derived from replicated logical state:

- local physical host/device allocation fails after deterministic eviction,
- local transfer submission fails,
- background IO reports partial/failure status,
- a debug/invariant check wants to prove the ranks are still aligned,
- shutdown/recovery needs to drain unknown pending state.

It should not be used merely to decide routine facts such as "which node was
evicted", "which prefix of write acks is logically complete", or "which host
range belongs to this request" if those facts can be represented by a shared
logical op id and deterministic local event progress.

Practical target:

- **Fast path:** no collective communication.
- **Normal async completion:** commit by deterministic op order and local event
  readiness, with zero-owned ranks carrying no-op events.
- **Slow path:** one coarse global success/abort gate after local retry has
  failed, not repeated per node or per scheduler tick.
- **Debug mode:** optional all-rank invariant checks, off the performance path.

This does not mean "never all-reduce". It means all-reduce is reserved for
cases where the program cannot otherwise know whether every rank is still able
to complete the same logical operation. If ordinary progress requires frequent
all-reduce, the local state machine is carrying insufficient information.

### Replicated capacity ledger

To avoid using all-reduce as an admission oracle, the scheduler should maintain
a replicated logical capacity ledger. The ledger is updated only by synchronous
logical ops, not by background transfer callbacks.

For CP shared KV, the ledger should track at least:

```text
device_free_pages_by_owner
device_evictable_pages_by_owner
host_free_pages_by_owner
host_evictable_pages_by_owner
pending_load_pages_by_owner
pending_backup_pages_by_owner
```

The owner dimension is necessary because a load-back allocation is not just
"N tokens"; it requires a specific owner-lane pattern. The host dimension is
necessary because CP host pools are physically local even though radix state is
logical and replicated.

With this ledger, common decisions become deterministic:

- whether a host-hit load can reserve device pages,
- whether a cold extend can reserve device pages,
- how many host victims are needed before backup,
- which logical host victims should be selected,
- whether a request should be admitted or kept queued.

The actual local allocator then becomes an executor of a decision already made
by the ledger. If the local allocator fails despite the ledger saying success,
that is either a bug, stale accounting, or an external failure. It should go to
slow-path recovery/debug validation, not become the normal admission mechanism.

### Visibility gates instead of confirmation votes

Async operations should become request-visible through deterministic gates:

```text
submit logical op in same order on every rank
reserve ledger capacity synchronously
submit local transfer or no-op event
later, when a request/victim decision needs the result:
  wait/query the local event for that op
  expose the next deterministic op state
```

This avoids asking every rank "are you done yet?" on every tick. A faster rank
may finish its local D2H/H2D earlier, but it should not make a different
logical scheduling decision. It can only expose the result at a deterministic
visibility point for the same logical op.

Examples:

- Backup is submitted as `HOST_BACKUP_PENDING`. It becomes `HOST_VALID` only
  at a backup visibility gate. Until then, host-hit and host-evict logic treats
  it as pending.
- Load is submitted as `DEVICE_LOAD_PENDING`. The owning batch may consume it
  through `hicache_consumer_index`; other requests cannot treat it as ordinary
  `DEVICE_VALID` until the load visibility gate commits.
- Host eviction selects the same logical victims from the replicated host leaf
  set and ledger. Each rank frees only its local host slots. There is no need
  to vote on the victim list.

The important restriction is that visibility gates must be tied to logical
operation order, not wall-clock completion. Local event completion is a data
dependency; logical visibility remains controlled by the radix/scheduler state
machine.

## MTP/Draft KV Pool Model

EAGLE/MTP draft KV should be treated as a shadow payload of the target KV
radix node, not as an independent cache.

The draft KV pool must not maintain its own:

- radix tree,
- hit/miss decision,
- backup/load/evict state,
- owner pattern,
- prefix length,
- protection state,
- logical commit/abort decision.

The target radix node is the source of truth. If MTP is enabled, target and
draft residency are a single atomic logical state:

```text
target HOST_VALID  <=> draft HOST_VALID
target DEVICE_VALID <=> draft DEVICE_VALID
```

Partial states such as target-host-valid but draft-host-absent are invalid for
request-visible cache hits, because decode acceptance depends on target and
draft KV representing the same prefix.

The target node metadata therefore owns both target and draft physical payloads:

```text
CpHiCacheNodeMetadata:
  logical_len
  owned_positions
  page_owners
  host_indices        # target host slots, local physical shard
  draft_host_indices  # draft host slots, local physical shard when MTP enabled
```

For zero-owned ranks, both `host_indices` and `draft_host_indices` are empty.
For non-zero-owned ranks under MTP, the target and draft host index arrays must
have the same owned-position cardinality.

Lifecycle implications:

- Backup reserves target and draft host slots together.
- Draft reservation or D2H submission failure rolls back target reservation.
- Load reuses the target node's owner pattern and owned positions for draft
  H2D.
- Device eviction frees target and draft device payload together.
- Host eviction frees target and draft host payload together.
- Radix split must split target and draft metadata together, or drain the
  pending op before splitting.

This avoids the earlier failure mode where cache-hit prefix target KV was
restored but draft/MTP KV was stale, absent, or read from a mismatched owner
layout, causing speculative accept length to collapse after HiCache hits.

## Eviction Scenarios In Scope

For the first implementation stage, only these eviction cases need to be
handled:

1. **Host-hit load-back needs device KV but device space is insufficient.**
   - Evict only device cache that is host-backed, unlocked, not under D2H
     backup, and not a load destination.
   - Do not affect the currently running prefill request.
   - Use owner-lane-aware capacity checks before committing load-back.

2. **Cold prefill extend needs device KV but device space is insufficient.**
   - Evict host-backed device cache first.
   - Use the same CP owner-lane rules as load-back.
   - If insufficient, reject/fallback at scheduling time rather than partially
     committing a batch.

3. **Write-through backup needs host KV but host space is insufficient.**
   - Evict host cache only from nodes that are host-valid, device-absent, not a
     load source, and not a backup destination.
   - Free target and draft host slots atomically.
   - Select the same logical victims across ranks; each rank frees only its
     local physical host slots.

4. **Low-frequency proactive host free-room maintenance.**
   - Optional after correctness is stable.
   - Should run in coarse epochs, not every scheduler tick.

Out of scope for this stage:

- decode prealloc and decode-side append pressure,
- TBO/two-batch overlap reserve accounting,
- L3/storage write-back or prefetch,
- async write-back demotion as a required emergency path,
- cross-instance cache sharing.

## Backup/Load Overlap Cases to Support

There are two distinct overlap opportunities. They should not be conflated
because their correctness constraints differ.

### Case A: computed suffix backup while the same request is still forwarding

For a new request or cache-miss suffix, each layer's KV becomes immutable once
that layer has stored its KV for the current prefill. In principle, that
layer's D2H backup can start before later layers finish computing:

```text
layer i stores target KV
  -> record layer-i store completion
  -> D2H backup layer i on write_stream
layer i+1 compute continues on compute stream
```

This can reduce backup tail latency and make host cache visible sooner after
the request completes. It does not allow device KV to be evicted earlier for
the same logical node unless the implementation supports partially host-valid
nodes, which is not recommended for the first stage.

Required conditions for per-layer backup:

- A layer completion event must order D2H after the layer's KV store.
- Target and draft/MTP backup must both participate in the same logical backup
  op.
- NSA indexer K/scale payload must be backed up with the corresponding layer,
  not forgotten behind the normal KV path.
- The node becomes request-visible `HOST_VALID` only after every required
  target and draft layer has completed and the global CP commit succeeds.
- Host eviction must not reclaim any layer slice of a pending backup
  destination.

The current host-pool API does not expose a generic
`backup_from_device_per_layer()` equivalent to `load_to_device_per_layer()`.
Adding per-layer backup therefore requires a host-pool API extension for MHA,
MLA, and NSA indexer payloads. Until then, all-layer backup should remain the
correctness baseline.

### Case B: host/storage prefetch while the request later backs up computed KV

If a request already has host-cache or storage-cache prefix, the system may
prefetch additional pages into host while the model recomputes or extends a
suffix. Later, the computed suffix may also be backed up to host. This creates
a conflict domain over both host memory and radix ownership:

```text
storage/L3 prefetch writes host slots for logical pages P
compute finishes and write-through backup also writes host slots for P
```

The safe rule is:

> Compute-derived backup wins over speculative prefetch for the same logical
> token range.

Reason: compute-derived KV is exactly the data used by this request's forward,
while prefetch is opportunistic and can be partial, revoked, or stale with
respect to the current radix split state.

Practical consequences:

- Prefetch must be request-local and non-authoritative until insertion commit.
- Backup insertion for an overlapping logical range should either cancel the
  overlapping prefetch or make prefetch insertion skip/free the already-backed
  range.
- Prefetch host slots and backup host slots must both participate in host
  memory pressure accounting.
- Host eviction must not choose active prefetch destinations or active backup
  destinations.
- Radix insertion remains synchronous; background prefetch threads can only
  deliver completed bytes and completion counts.

Current code already has pieces of this model: prefetch protects the anchor
host node, records `ongoing_prefetch`, and `_insert_helper_host()` skips matched
prefix ranges when inserting prefetched host data. For CP shared KV + MTP,
however, the same principle needs to include target/draft atomicity and owner
metadata, not just host byte insertion.

## Current Risk Areas

### 1. Backup lacks an explicit global success commit gate

`write_backup()` sets `node.host_len` and `node.cp_hicache` after local success.
If one rank succeeds and another fails after retry, logical host-backed state
can diverge unless all ranks reach the same outcome through surrounding control
flow.

Current host eviction retry improves this, but it is not a complete transaction
protocol. A stronger model should separate:

```text
local reserve success -> global success check -> logical metadata commit
```

### 2. Load failure can become rank-local

`load_cp()` returns `None` if owner-lane allocation fails. `load_back()` retries
with lane-aware eviction, then returns `None` and falls back to recompute if it
still fails.

The expected case is that all ranks make the same decision because the logical
allocator is replicated. Under pressure or bugs, this assumption can break.
The slow path should use a global abort/commit gate after retry.

### 3. Host eviction synchronization is too fine-grained

`_evict_host_for_physical_slots()` can all-reduce inside the victim loop. This
is correct but expensive when host memory is close to full and many small
evictions happen.

A batched epoch model should reduce synchronization frequency.

### 4. Write-back eviction is blocking

When `write_policy == "write_back"`, current device eviction waits for D2H
completion before freeing device slots. This is safe but does not provide the
performance benefit expected from asynchronous demotion.

Async demotion is possible, but it cannot help an immediate allocation failure
unless the system can tolerate delayed device release. Emergency eviction may
still need to drop cache instead of waiting.

### 5. Write-through can create excessive ack all-reduce pressure

`write_through` uses threshold 1, so every inserted node can trigger backup.
`writing_check()` then may all-reduce frequently even if only a small number of
acks have completed. This can become a scheduler hot path.

### 6. Radix node split can race with async node state

Radix tree maintenance is not just metadata bookkeeping under async HiCache.
The tree can split nodes during `match_prefix()` or `insert()`:

- `HiRadixCache.match_prefix()` may call `_match_prefix_helper()`, which can
  call `_split_node()` on a partial match.
- `HiRadixCache.insert()` can also call `_split_node()` on a partial match.
- `_split_node()` copies `child.lock_ref` into the newly inserted parent node
  and splits `value`, `cp_hicache`, `host_len`, and hash metadata when present.

This is correct for quiescent nodes, but it is risky for nodes with in-flight
async operations.

Example risk:

1. A new node is inserted.
2. `write_through` starts async backup for that node.
3. `write_backup()` records `ongoing_write_through[child.id] = child` and uses
   `inc_node_lock_ref(child)`.
4. Before the D2H ack commits, another request partially matches the node and
   `_split_node()` inserts a new parent.
5. The new parent inherits `lock_ref`, and CP host metadata is split between
   new parent and child.
6. The write ack still only references the original child id.
7. `writing_check()` releases the child lock, but has no ack entry for the new
   parent lock.

This can leave the split parent permanently protected or otherwise make
`evictable_size_`, `protected_size_`, and leaf sets diverge from actual radix
state. A path lock is less exposed because `dec_lock_ref(child)` walks through
the inserted parent after split, but the write-through single-node lock path is
specifically vulnerable unless split is made op-aware.

This suggests a general rule:

> A radix node with in-flight async state must either be unsplittable until the
> operation commits, or the async operation must be explicitly split/aliased
> together with the radix node.

### 7. GPU memory admission tracks KV tokens, not all transient memory

Current scheduling largely reasons in token-pool units:

- `PrefillAdder.rem_total_tokens` uses allocator `available_size()` plus radix
  `evictable_size()`, then subtracts running/new request offsets.
- `_get_load_back_mem_quota()` reserves `real_input_tokens + page_size` before
  deciding how much host cache can be loaded back.
- `alloc_paged_token_slots_extend()` evicts by token count before allocating
  extend KV slots.
- CP owner-lane retry uses `_evict_for_compute_owner_lanes()`.

This protects the KV allocator from overcommit, but it is not a full CUDA memory
budget. It does not directly account for:

- model forward temporary buffers,
- logits / hidden-state capture,
- MLP / DeepGEMM / EP communication scratch,
- CUDA graph memory,
- two-batch overlap staging,
- multiple concurrent HiCache load producers,
- draft/MTP compute-side temporary memory.

The model runner profiles KV token capacity from `mem_fraction_static`, then CP
shared KV expands logical capacity by `attn_cp_size` while keeping the physical
pool at the profiled size. That means async HiCache must leave explicit
headroom at the scheduler layer for non-KV transient memory; otherwise token
accounting can look valid while CUDA still OOMs in a later model kernel.

## Recommended State Model

Introduce a logical async operation journal for CP HiCache. The journal can be
implemented incrementally around current queues, but the state model should be
explicit.

Each operation should contain:

```text
op_id
kind: BACKUP | LOAD | DEVICE_EVICT | HOST_EVICT
node_ids
logical_len
page_owners hash / metadata
target_physical_count
draft_physical_count
zero_owned flag
state: PREPARED | SUBMITTED | LOCAL_DONE | GLOBAL_COMMITTED | ABORTED
stream_event
```

Rules:

1. All ranks append the same logical ops in the same order.
2. Physical transfer can be zero length.
3. Commit only advances a contiguous op prefix.
4. Commit is logical; transfer completion is local.
5. Rollback is local physical cleanup plus global logical abort.
6. Radix split/merge must update or block pending ops that reference affected
   nodes.

## Proposed Synchronization Policy

### Fast path: deterministic, no collective sync

Use only local events and existing streams for:

- D2H backup submission.
- H2D load submission.
- Per-layer load/compute overlap.
- Zero-owned no-op transfer.
- Device eviction of nodes already host-backed.

Fast-path correctness comes from logical op ids, radix order, and owner-pattern
metadata. A rank does not ask other ranks whether a page belongs to it; it
derives that from `page_owners`. A rank does not ask which node is being
committed; it consumes the next logical op in the same queue.

### Normal commit path: local-event prefix commit

For async backup/load completion, the desired normal path is:

```text
local event for op N is complete
zero-owned ranks use an already-complete event
logical op queue is consumed in order
```

The commit operation should advance only a contiguous local completed prefix.
If all ranks have the same logical op order and every op has a local event,
then each rank can commit the same semantic prefix without per-tick global
voting. If a rank's local event is not ready, that rank simply does not expose
the later logical state yet; request scheduling must respect its local pending
state.

### Batched validation path: amortized all-reduce

Use `all_reduce(MIN)` on a monotonically increasing completed-op prefix only
when useful:

- Debug/invariant validation is enabled.
- Local completed prefix has advanced by a large threshold and the system wants
  to assert cross-rank alignment.
- Ongoing write count exceeds a safety threshold.
- Scheduler is about to make a memory-sensitive admission decision.
- Memory pressure requires releasing write locks.
- Idle checks need to drain async ops.

Avoid all-reduce on every scheduler tick when no local progress happened.

### Slow path: global success/abort gate

Use collective sync only after local retry paths:

- Host reserve failed and host eviction/retry was attempted.
- Draft host reserve failed.
- Owner-lane device allocation failed and lane-aware eviction/retry was
  attempted.
- Write-back demotion cannot reserve host memory.

Protocol:

```text
local_success = reserve/alloc/submit success
global_success = all_reduce_min(local_success)
if global_success:
    commit logical op
else:
    rollback local physical reservations
    abort logical op identically on all ranks
```

## Proposed Backup Lifecycle

```text
RESIDENT
  -> BACKUP_PREPARED      reserve target/draft host slots locally
  -> BACKUP_SUBMITTED     enqueue D2H on write_stream
  -> BACKUP_LOCAL_DONE    finish_event ready locally
  -> HOST_VALID           global contiguous commit releases device write lock
```

Failure handling:

- Host reserve failure enters host eviction slow path.
- Draft reserve failure must roll back target reserve.
- Global failure leaves node device-resident but not host-backed.
- Write-through failure should skip backup rather than block request progress.
- Write-back failure should not assert; it needs a clear fallback path.

Request-visible state rule:

- `BACKUP_PREPARED` and `BACKUP_SUBMITTED` must not be treated as
  `HOST_VALID`.
- The device source remains protected until the global contiguous backup commit
  releases the write lock.
- Host eviction must not select backup destinations that have not globally
  committed.

For the first implementation, this can be represented without a large new
journal by adding an explicit "pending host backup" state around the existing
`ongoing_write_through` and `ack_write_queue` path. The important behavioral
change is to distinguish:

```text
host slots reserved and D2H submitted
```

from:

```text
host bytes are valid and request-visible
```

## Proposed Load Lifecycle

```text
HOST_VALID + DEVICE_EVICTED
  -> LOAD_PREPARED        alloc_pages_with_owners(page_owners)
  -> LOAD_SUBMITTED       enqueue per-layer H2D on load_stream
  -> LAYER_READY[i]       LayerDoneCounter event for layer i
  -> DEVICE_VALID         final load ack releases load lock
```

Failure handling:

- Owner-lane allocation failure enters lane-aware eviction slow path.
- After retry, global success/abort gate decides whether all ranks load or all
  ranks recompute.
- If aborted, radix host state remains valid, but no device state is committed.

Request-visible state rule:

- `LOAD_PREPARED` and `LOAD_SUBMITTED` are valid only for the batch that owns
  the returned `hicache_consumer_index`.
- Other requests should not treat the loaded node as generally
  `DEVICE_VALID` until the load ack commits.
- Host metadata remains valid throughout load; aborting a load should fall back
  to recompute without invalidating the host cache.

Current code already has a per-layer data dependency through
`LayerDoneCounter`, which is the right data-plane primitive. The missing
control-plane distinction is that `node.value` is assigned before the H2D load
fully completes. That is acceptable only if the scheduler ensures every
consumer of those device locs carries the same `hicache_consumer_index`, or if
the node remains in a `DEVICE_LOAD_PENDING` logical state until the load ack
commits.

For the first implementation, prefer the second rule:

```text
load owner batch may consume DEVICE_LOAD_PENDING with producer id
all other requests see the node as host-valid but not device-valid
```

This keeps layer-wise overlap for the triggering batch while preventing later
requests from racing on partially loaded device KV.

## Proposed Eviction Lifecycle

### Device eviction

Host-backed node:

```text
DEVICE_VALID + HOST_VALID
  -> DEVICE_EVICTED + HOST_VALID
```

This remains cheap and synchronous in metadata because no transfer is required.

Node under backup:

```text
BACKUP_SUBMITTED
  -> not evictable until backup commit or abort
```

Write-back demotion:

```text
DEVICE_VALID + HOST_INVALID
  -> DEMOTE_SUBMITTED
  -> HOST_VALID + DEVICE_EVICTED
```

This can be async for proactive demotion. For emergency allocation, the system
may need to drop cache instead of waiting for demotion to finish.

### Host eviction

Use batched epochs:

```text
collect victims by logical priority
free local physical target/draft host slots
accumulate local_freed_slots
all_reduce_min(done) only after a batch/epoch
commit logical host removal for the common evicted prefix
```

This should replace per-node or per-iteration all-reduce in the common
host-pressure path.

## Radix Tree Async Safety Model

The radix tree should treat async backup/load/evict state as part of node
residency, not as external controller-only state.

Recommended per-node fields conceptually:

```text
device_state: ABSENT | VALID | LOADING | WRITING_SOURCE | DEMOTE_PENDING
host_state: ABSENT | WRITING | VALID | EVICTING
pending_ops: set/op refs
split_epoch: integer
```

The implementation does not need exactly these fields, but it needs equivalent
semantics.

### Safe split options

There are three viable designs.

#### Option A: Drain-before-split

Before `_split_node()` splits a node with pending async state, force the related
op to commit or abort.

Pros:

- Smallest correctness surface.
- Easy to reason about.
- Good first implementation.

Cons:

- Partial-prefix traffic can block on D2H/H2D completion.
- May reduce write-through/load overlap in workloads with frequent shared
  prefix refinement.

#### Option B: Op-aware split

When `_split_node()` splits a node with pending backup/load state, split the
pending op's logical segment as well.

Requirements:

- The ack must release all locks created by the original op after split.
- Host metadata must be split in the same way as radix metadata.
- The journal must understand that one submitted physical transfer corresponds
  to multiple post-split logical nodes.
- Commit must be idempotent if a later split happens before the ack.

Pros:

- Preserves maximum async overlap.

Cons:

- Highest complexity.
- Easy to get lock/ref accounting wrong.

#### Option C: Pending segment alias

Represent an async operation as an immutable logical segment object independent
of the current radix node object. A split creates two radix nodes that both
reference slices of the same pending segment. The op ack commits the segment,
then materializes per-node state through the current aliases.

Pros:

- Cleaner than mutating ack lists on every split.
- Good long-term model if async operations become more pervasive.

Cons:

- Requires a larger refactor around `ongoing_write_through`,
  `ongoing_load_back`, and split metadata.

Recommended path:

1. Start with **Option A** for backup/write-through nodes to eliminate the
   highest-risk lock leak.
2. Keep path-lock load behavior as-is initially, but add assertions/tests around
   splitting loaded/in-flight nodes.
3. Move toward **Option C** only if drain-before-split shows measurable
   performance loss.

### Leaf set and accounting rules

Every async state transition must update the same accounting surfaces:

- `evictable_size_`
- `protected_size_`
- `evictable_leaves`
- `evictable_host_leaves`
- `lock_ref`
- `host_ref_counter`

The recent class of bugs where `evictable_size_` claimed capacity but
`evictable_leaves` had no usable victims shows that size counters and leaf sets
must be updated atomically with node residency changes. Async journal commit
should own these updates rather than scattering them across controller callbacks
and radix helpers.

## GPU Memory Margin Policy

Token-pool accounting should be treated as necessary but not sufficient.

Recommended device memory budget layers:

1. **Hard KV allocator budget**
   - Existing `available_size + evictable_size - offsets`.
   - Must remain the source of truth for KV slot allocation.

2. **Owner-lane budget**
   - CP shared KV must check not only total logical tokens but also per-owner
     lane free/release pages.
   - Load-back and extend should both use owner-lane retry before falling back
     or aborting.

3. **Async in-flight reserve**
   - Loaded host hits occupy device KV slots before forward consumes them.
   - Write locks reduce evictable KV slots until D2H commits.
   - Future two-batch overlap must reserve for all active load producers and
     next-batch extend slots, not just the current batch.

4. **Non-KV CUDA headroom**
   - Reserve a fixed or profiled margin for logits, hidden states, MLP/EP
     scratch, CUDA graph memory, and speculative decoding temporaries.
   - This margin should be outside radix token accounting.

5. **Emergency reserve**
   - Keep at least one page per owner lane, or a small per-lane percentage, as
     an emergency reserve for forward progress.
   - Avoid consuming the last owner-lane pages for opportunistic load-back.

Practical admission rule:

```text
allow load-back only if:
  total_available_plus_evictable
    - real_extend_tokens
    - async_inflight_reserved_tokens
    - emergency_owner_lane_reserve
    >= host_hit_tokens

and each required owner lane has enough free/releasable pages after reserve.
```

For two-batch overlap, the scheduler should reserve the maximum of:

- current batch real extend tokens,
- current batch host load-back tokens,
- next prepared batch real extend tokens,
- next prepared batch host load-back tokens,
- per-slot logits/hidden temporary budget if those tensors are retained across
  overlap boundaries.

This is intentionally more conservative than pure token accounting. It is the
price of avoiding late CUDA OOM after the scheduler has already committed a
batch.

## Implementation Direction

Recommended order:

1. **Correctness gates first**
   - Backup global success/rollback after reserve retry.
   - Load global abort after owner-lane retry.
   - Target/draft atomic reserve and rollback.
   - Drain or make split-aware any radix node with in-flight backup state.

2. **Reduce synchronization frequency**
   - Batch write ack commits by completed prefix.
   - Skip all-reduce when local completed prefix has not changed.
   - Batch host eviction into epochs.
   - Move host eviction synchronization from per-victim to per-epoch.

3. **Then consider async write-back demotion**
   - Keep write-through and write-through-selective stable first.
   - Add write-back async demotion only with explicit emergency fallback.

4. **Add explicit memory reserves**
   - Owner-lane emergency reserve.
   - In-flight load/write reserve.
   - Non-KV CUDA headroom for overlap/speculative paths.

## Open Questions

1. What is the acceptable maximum delay before write locks are released under
   write-through? This determines ack batching thresholds.
2. Should host eviction be driven only on allocation failure, or should it keep a
   low-frequency free-room background target?
3. For write-back, should emergency eviction drop cache immediately when host
   reserve fails, or block until demotion completes?
4. Should the op journal be a new abstraction, or should existing
   `ack_write_queue`, `ack_load_queue`, `ongoing_write_through`, and
   `ongoing_load_back` be tightened first?
5. Should `_split_node()` drain pending backup ops, or should pending ops become
   split-aware?
6. How much non-KV CUDA headroom is required for GLM5 + EAGLE/MTP + EP under
   two-batch overlap?

## Summary

The current code already has the right low-level pieces:

- rank-local CP physical filtering,
- zero-owned no-op acks,
- async write/load streams,
- per-layer load/compute overlap,
- owner-pattern-preserving reload.

The missing piece is a unified logical transaction model. Without it, making
backup/load/evict more asynchronous increases the risk of cross-rank divergence.
The desired direction is:

> Fast path uses local async events only; commit path uses batched prefix
> all-reduce; failure path uses explicit global success/abort; physical pool
> operations remain rank-local.

Additionally, radix tree mutation and GPU memory admission must be part of the
same design. Async state cannot be treated as controller-private state because
radix split, lock release, leaf eligibility, and memory quota decisions all
depend on it.
