# NSA Prefill CP HiCache Load and Shared-KV Prefetch Overlap Notes

Date: 2026-05-28

Branch context: `cjy-cp-refactor` after the direct/page_first_direct HiCache
load/backup work.

This note records the current understanding of cache-hit prefill when all of the
following are enabled:

```text
--enable-nsa-prefill-context-parallel
--nsa-prefill-cp-mode in-seq-split
--enable-nsa-prefill-cp-shared-kv
--enable-hierarchical-cache
--hicache-io-backend direct
--hicache-mem-layout page_first_direct
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1
```

The important conclusion is that **HiCache L2->L1 load overlap is necessary but
not sufficient**.  In CP shared-KV mode, attention cannot consume the restored L1
KV shard directly.  It must still run the shared-KV compatibility path that
materializes dense prefix buffers and performs CP collectives.  Therefore the
cache-hit critical path has two distinct prefetch/load stages:

1. **HiCache host L2 -> device L1 shard load**.
2. **CP shared-KV L1 shard -> attention-visible dense prefix materialization**.

Optimizing only stage 1 can leave the real stall in stage 2.

---

## 1. Current cache-hit pipeline

### 1.1 L2 -> L1 HiCache load

On a host-cache hit, radix/scheduler logic first restores host-backed nodes into
fresh device KV slots:

```text
schedule_policy.py
  AddReq -> tree_cache.init_load_back(...)

hiradix_cache.py
  HiRadixCache.load_back(...)
    -> cache_controller.load_cp(...)

cache_controller.py
  HiCacheController.load_cp(...)
    -> alloc_pages_with_owners(page_owners)
    -> enqueue target/draft load ops

scheduler.py
  ready_to_load_host_cache()
    -> cache_controller.start_loading()
```

For CP shared KV, `load_cp()` preserves the write-time owner pattern via
`node.cp_hicache.page_owners`.  This is required because a later cache hit
allocates fresh logical pages; if the owner sequence changes, each rank would load
its host bytes into slots interpreted as another rank's shard.

With `direct/page_first_direct`, each layer uses the TAI direct H2D path:

```text
memory_pool_host.py
  load_to_device_per_layer(..., io_backend="direct")
    -> tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf(...)
```

The TAI op uses `cudaMemcpyBatchAsync`, so this path does not spend SM cycles on
an H2D copy kernel when the runtime capability is available.  It is submitted on
SGLang's `load_stream`.

### 1.2 Layer-wise wait point

`start_loading()` submits per-layer H2D work and records one layer event per
layer.  The forward path receives `hicache_consumer_index`, and the KV pool waits
when a layer's buffer is actually accessed:

```text
memory_pool.py
  get_key_buffer(layer_id) / get_value_buffer(layer_id)
    -> layer_transfer_counter.wait_until(layer_id)
```

So the current L2->L1 behavior is:

```text
load_stream:     load L0 -> event0 -> load L1 -> event1 -> load L2 -> ...
forward_stream:  wait L0 -> compute L0 -> wait L1 -> compute L1 -> ...
```

This provides layer-wise overlap, but it is not a complete cache-hit solution for
shared KV because restored L1 KV remains sharded.

### 1.3 CP shared-KV dense prefix materialization

In CP shared-KV mode, persistent KV/index cache is sharded by owner lane.  Before
MLA/NSA attention can use a historical prefix, each layer still needs to build an
attention-visible dense view:

```text
shared physical L1 shards
  -> local materialize/remap into dense prefix buffer
  -> CP all-reduce / range reduce
  -> attention consumes dense rows/pages
```

Phase 8 adds one-layer-ahead prefix prefetch for two payloads:

- MLA KV prefix.
- NSA index K/scale prefix.

The intended Phase 8 flow is:

```text
Layer L prepare:
  start Layer L+1 prefix index + MLA materialize prefetch

Layer L attention:
  consume Layer L prefetched prefix if available
  synchronously materialize/reduce current suffix only
  run attention

Layer L+1 attention:
  wait for Layer L+1 prefix prefetch only if it is not done yet
  synchronously materialize/reduce current suffix only
```

This means a cache-hit layer may have both:

```text
HiCache L2->L1 H2D readiness wait
CP shared-KV prefix materialization/reduce readiness wait
```

A useful cache-hit overlap design must reason about both.

---

## 2. Why the current overlap is weaker than expected

The current direct/page_first_direct HiCache load path overlaps only the physical
host-to-device restore.  It does not by itself hide the shared-KV compatibility
work that follows.

### 2.1 Front-loaded L2->L1 submit

`start_loading()` is called before `prepare_for_extend()`/forward execution and
submits per-layer load ops in a Python loop.  The copy work is asynchronous, but
submission and index preparation are front-loaded before the model starts doing
layer compute.

This is acceptable for correctness, but it means L2->L1 overlap is not the same
as "submit exactly at layer boundary".  The first layer still waits, and CPU
submit overhead is paid before useful model compute can hide it.

### 2.2 Shared-KV prefetch depends on restored L1 KV buffers

The Phase 8 prefetcher eventually reads `token_to_kv_pool.get_key_buffer(...)` or
`get_index_k_with_scale_buffer(...)` for the next layer.  In a HiCache hit, those
buffer getters are also guarded by the HiCache layer-transfer counter.  Therefore
the CP shared-KV prefetch launch can be implicitly delayed by the L2->L1 layer
load event.

Effective dependency:

```text
L2 host page -> L1 physical shard loaded
  -> shared-KV prefix prefetch can read that layer's shard
  -> dense prefix materialize + CP reduce
  -> attention can consume prefix
```

If L2->L1 load for layer `L+1` is late, starting the shared-KV prefetch for
layer `L+1` from layer `L` may still block at the buffer getter and lose the
intended overlap window.

### 2.3 Prefix prefetch and HiCache load use different readiness abstractions

HiCache load readiness is represented by `LayerDoneCounter` events attached to
KV pool access.  Shared-KV prefetch readiness is represented by prefetch handles
and CUDA events inside `CpSharedKVMlaPrefetcher` and
`CpSharedKVIndexPrefetcher`.

These mechanisms are both correct locally, but they are not currently planned as
a single cache-hit pipeline.  As a result, each stage can appear overlapped in
isolation while the composed pipeline still serializes at consumption points.

### 2.4 Bandwidth and stream contention remain possible

The direct/page_first_direct L2->L1 load uses copy-engine H2D, while shared-KV
materialization/reduce uses GPU work and NCCL/PyNccl collectives.  They are not
the same operation, but in production they can still contend for:

- host pinned-memory bandwidth,
- PCIe/NVLink copy resources,
- NVLink/NCCL collective bandwidth,
- stream scheduling and event dependencies,
- CPU submission time.

Seeing no end-to-end cache-hit speedup after L2->L1 load overlap does not prove
that L2->L1 overlap is broken; it can mean the post-load shared-KV prefix path is
now the limiting stage.

---

## 3. Improved overlap model

For CP shared KV + HiCache, the useful mental model is a two-stage prefetch DAG,
not a single H2D load:

```text
For layer L+1 on a cache-hit request:

Stage A: HiCache restore
  host L2 page_first_direct target/draft/index payload
    -> direct H2D into owner-lane L1 physical slots
    -> layer-load event A[L+1]

Stage B: shared-KV prefix prefetch
  wait/read L1 physical shard for L+1 prefix
    -> local dense prefix materialize/remap
    -> async CP reduce / prefix event B[L+1]

Stage C: layer L+1 attention consume
  wait A[L+1] only when raw L1 buffer is needed
  wait B[L+1] only when dense prefix is needed
  materialize current/suffix rows synchronously
  run attention
```

The target is to launch Stage A and Stage B early enough that Stage C rarely
waits on either event.

### Desired staggered timeline

The expected cache-hit pipeline is **two-layer-ahead L2->L1 prefetch** plus the
existing **one-layer-ahead shared-KV L1 prefix prefetch**:

```text
Layer L-2 / batch setup:
  start L2->L1 HiCache load for Layer L
  target event: raw owner-lane L1 shard for Layer L becomes ready before Layer L-1

Layer L-1 prepare:
  start Layer L shared-KV prefix prefetch from the restored L1 shard
  prefetch_stream waits on Layer L L2->L1 load event if needed
  forward_stream should not block just to launch this prefetch

Layer L-1 attention / MLP:
  hide Layer L shared-KV prefix materialize and CP reduce

Layer L attention:
  consume prefetched prefix
  only do current/suffix materialize on the critical path
```

Equivalently, for a layer `X` that will consume cache-hit prefix data:

```text
Stage A: L2->L1 load for X        starts at X-2 or earlier
Stage B: L1 shared-KV prefetch X  starts at X-1
Stage C: attention consume X      runs at X
```

The two-layer lead for Stage A is not an arbitrary latency knob.  It follows from
the dependency chain: Stage B reads L1 physical shards, so Stage A must already be
far enough ahead that Stage B can launch without stealing the current layer's
compute window.

### Practical implication

`SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` should be evaluated together with
HiCache load state.  A prefetcher that starts "one layer ahead" but immediately
waits for the L2->L1 load of that next layer has less overlap than expected.

The current all-layer `start_loading()` submission is not by itself proof that
Stage A has the required two-layer lead.  The useful property is readiness at the
Stage-B launch point: by Layer `L-1` prepare, Layer `L`'s L2->L1 load should be
complete or represented as a local CUDA event that the prefetch stream can wait
on without blocking the forward stream.

---

## 4. Concrete improvement directions

### 4.1 Make cache-hit prefetch dependency explicit

Do not treat HiCache load and shared-KV prefetch as independent optimizations.
Represent their dependency explicitly in the design:

```text
shared-KV prefix prefetch for layer X requires L2->L1 load event for layer X
```

This does not necessarily require a new collective.  It can be expressed through
local CUDA events because the dependency is local to each rank's physical shard.

### 4.2 Start L2->L1 load early enough for next-layer shared-KV prefetch

The current load stream submits all layers before forward, which is already early
for the copy itself.  The remaining question is whether CPU submission and copy
engine progress are far enough ahead of Phase 8's next-layer prefetch hook.

The useful measurement is not just "load submitted".  It is:

```text
At Layer L-1 prepare, is L2->L1 load for Layer L already complete or close enough
that shared-KV prefix prefetch can launch without stalling?
```

### 4.3 Chain shared-KV prefix prefetch after load events without blocking the
current layer

If the next-layer L1 shard is not ready yet, the prefetch path should record a
stream wait on the load event and enqueue the prefetch work on the prefetch
stream, rather than blocking the current forward stream while calling a KV buffer
getter.

Conceptually, when Layer `L-1` wants to prefetch Layer `L`:

```text
prefetch_stream.wait_event(load_event[L])
prefetch_stream: materialize prefix for L + async reduce
forward_stream: continue Layer L-1 compute
```

This preserves the dependency while keeping the current layer's compute window
available for overlap.

A direct call path that goes through `get_key_buffer(L)` or
`get_index_k_with_scale_buffer(L)` can accidentally turn this into a forward
stream wait, because those accessors enforce `LayerDoneCounter.wait_until(L)`.
The improved design therefore needs a non-blocking way for the shared-KV
prefetcher to depend on the HiCache load event, or a raw-buffer access path used
only after the prefetch stream has been ordered behind that event.

### 4.4 Keep suffix/current materialization separate

Prefix is safe to prefetch because it already exists in the persistent pool after
cache hit/load.  Current/suffix pages for layer `L+1` are not ready until layer
`L+1` prepare writes them.  The existing Phase 8 split remains correct:

```text
prefix: prefetchable
current/suffix: materialize synchronously at consume time
```

Do not try to prefetch current/suffix before the layer has written its local KV.

### 4.5 Target and draft remain one logical cache object

The L2->L1 load dependency applies to target and draft/MTP payloads.  Shared-KV
prefix materialization currently concerns target attention buffers, but cache-hit
correctness and speculative accept rate still require draft KV residency to move
with target KV.  Any improvement must preserve:

```text
target load/backup/evict visibility == draft load/backup/evict visibility
```

---

## 5. What to measure next

A benchmark or trace should separate the following timings per layer:

1. L2->L1 H2D load submit time.
2. L2->L1 H2D load completion event time.
3. Shared-KV prefix prefetch launch time.
4. Shared-KV prefix reduce completion time.
5. Layer attention consume wait time.
6. Current/suffix materialize/reduce time.

The key diagnostic questions are:

- At `Layer L-1` prepare, is `load_event[L]` already complete, or can the
  prefetch stream wait on it without blocking the forward stream?
- Did the L2->L1 load for `Layer L` start at `Layer L-2` or earlier in the
  effective timeline, not merely get enqueued somewhere before forward?
- If the load is not complete, does the current forward stream block while trying
  to start `Layer L` shared-KV prefetch?
- At `Layer L` attention, is the prefix prefetch event complete?
- Is the remaining stall mostly L2->L1 H2D, prefix materialize/reduce, or suffix
  materialize/reduce?
- Does enabling `direct/page_first_direct` reduce H2D time but expose the
  shared-KV prefix reduce as the new bottleneck?

Without these separated timings, an end-to-end "cache hit has no speedup" result
is ambiguous.

---

## 6. Current confidence and limits

High-confidence facts from code/docs:

- HiCache CP load is per-layer and uses `LayerDoneCounter` for layer-wise waits.
- `direct/page_first_direct` load routes to TAI `cudaMemcpyBatchAsync` H2D.
- Phase 8 shared-KV prefetch is one-layer-ahead and handles both MLA KV prefix
  and NSA index prefix.
- Shared-KV prefetch consumes L1 KV/index buffers and then performs dense
  materialization plus CP reduce.

Medium-confidence inference:

- In cache-hit workloads, Stage B can dominate after Stage A is optimized.  This
  explains why improving L2->L1 load alone may not improve end-to-end latency.

Unknowns requiring measurement:

- How often next-layer shared-KV prefetch blocks on the HiCache layer load event.
- Whether the main stall is CPU submit, H2D bandwidth, NCCL reduce, materialize
  local copy, or suffix/current path.
- Whether current production traces show copy/collective contention or simply a
  missing/late prefetch launch.


---

## 7. Proposed implementation design: make Stage-B prefetch wait on Stage-A locally

### 7.1 Current concrete conflict points

The conflict is not that L2->L1 and L1 shared-KV prefetch are conceptually
incompatible.  The conflict is that the current prefetch code enters the same
synchronizing accessors used by the attention consume path.

Current blocking points:

```text
CpSharedKVMlaPrefetcher.maybe_create(...)
  token_to_kv_pool.get_key_buffer(first_layer_id)

CpSharedKVMlaPrefetcher.start_next_layer_prefix(...)
  token_to_kv_pool.get_key_buffer(next_layer_id)

CpSharedKVIndexPrefetcher.maybe_create(...)
  token_to_kv_pool.get_index_k_with_scale_buffer(first_layer_id)

CpSharedKVIndexPrefetcher.start_next_layer_prefix(...)
  token_to_kv_pool.get_index_k_with_scale_buffer(next_layer_id)
```

Those getters intentionally wait on HiCache layer-load readiness:

```text
get_key_buffer(layer_id)
  -> layer_transfer_counter.wait_until(layer_id - start_layer)

get_index_k_with_scale_buffer(layer_id)
  -> layer_transfer_counter.wait_until(layer_id - start_layer)
```

That is correct for the final consumer.  It is wrong for prefetch launch when the
call is made from the forward path, because the wait is attached to the current
stream instead of the prefetch stream.

A second issue is that `start_next_layer_prefix()` currently runs local prefix
materialization before creating the pending handle.  Because it is called from
layer-forward hooks, this materialization also tends to execute on the current
forward stream instead of the shared-KV prefetch stream.

Net effect:

```text
Layer L forward hook wants to prefetch Layer L+1
  -> calls synchronizing getter for L+1
  -> forward stream waits for L2->L1 load_event[L+1]
  -> local prefix materialization runs on forward stream
  -> async reduce is only launched after that
```

The desired effect is:

```text
Layer L forward hook wants to prefetch Layer L+1
  -> prefetch stream waits for L2->L1 load_event[L+1]
  -> local prefix materialization runs on prefetch stream
  -> prefix reduce runs on the same FIFO prefetch stream
  -> forward stream continues Layer L work
```

### 7.2 Design invariant

For cache-hit CP shared KV, the Stage-B shared-KV prefetch for layer `X` must have
an explicit local dependency on the Stage-A HiCache L2->L1 load event for layer
`X`.

Stage-A must also finish **capacity reservation in CP owner-lane vector form**
before any Stage-B prefetch is considered valid.  In other words, the prefetch
repair is not only a stream/event repair:

```text
old capacity model:
  total_available_tokens >= total_required_tokens

new CP shared-KV capacity model:
  for each owner lane o:
    available_pages_by_owner[o] >= required_pages_by_owner[o]
```

The shared-KV prefix prefetcher must never discover L1 capacity shortage.  If it
does, the bug is in Stage-A reservation/eviction, not in Stage-B prefetch.

This must be local CUDA stream/event ordering only:

- no new CP all-reduce,
- no cross-rank polling,
- no CPU synchronization,
- no silent fallback to sync full materialization.

The only ordering needed is local because each rank only reads its own restored
owner-lane shard before the existing shared-KV collective materializes the dense
prefix.

### 7.3 Narrow API addition

Add a narrow prefetch-only KV-pool API instead of weakening the normal getters.
The normal getters should continue to block because attention consume must remain
safe.

Suggested shape:

```python
def wait_layer_transfer_on_stream(self, layer_id: int, stream) -> None:
    """If a HiCache layer transfer is active, make `stream` wait for this layer.

    This must not wait on the caller/current stream unless `stream` is the current
    stream by construction.
    """


def get_key_buffer_for_prefetch(self, layer_id: int, stream) -> torch.Tensor:
    self.wait_layer_transfer_on_stream(layer_id, stream)
    return self._get_key_buffer_no_transfer_wait(layer_id)


def get_index_k_with_scale_buffer_for_prefetch(self, layer_id: int, stream) -> torch.Tensor:
    self.wait_layer_transfer_on_stream(layer_id, stream)
    return self._get_index_k_with_scale_buffer_no_transfer_wait(layer_id)
```

Implementation notes:

- `LayerDoneCounter` already owns per-layer events through
  `events[consumer_index].load_events[layer_index]`.
- Add a method that records `stream.wait_event(...)` for every active consumer
  index, mirroring `wait_until()` but using the supplied stream.
- For layers outside pool range, preserve existing `layer_in_pool` checks.
- If no `layer_transfer_counter` exists, return raw buffers directly.
- Do not expose raw getters broadly.  Keep them private or name them
  `*_for_prefetch` to prevent accidental unsafe attention use.

For MLA pools where `_get_key_buffer()` does not already exist, add a private raw
helper and keep public `get_key_buffer()` behavior unchanged.

For NSA state/index pools, add a raw helper for
`index_k_with_scale_buffer[layer_id - start_layer]` and route only prefetch
through it.

### 7.4 Prefetcher changes

For both MLA and index prefetchers:

1. Stop using synchronizing getters in `maybe_create()` and
   `start_next_layer_prefix()`.
2. Use the prefetch-safe getter with the prefetcher's FIFO stream.
3. Move local prefix materialization into `with torch.cuda.stream(self.stream)`.
4. Keep the existing single FIFO stream shared by index and MLA prefix prefetch.
   This avoids rank skew across PyNccl collectives.
5. Continue to create one handle per layer and consume it with the existing
   `handle.event` wait at attention time.

The new `start_next_layer_prefix()` shape should be:

```text
start_next_layer_prefix(next_layer_id):
  if disabled/already_started/out_of_pool: return

  kv_or_index_buffer = pool.get_*_for_prefetch(next_layer_id, self.stream)

  with prefetch_stream:
    # prefetch_stream is already ordered after load_event[next_layer_id]
    dense_buffer = allocate dense prefix buffer
    materialize local prefix pages/rows into dense_buffer
    event = async CP reduce on dense prefix range

  handles[next_layer_id] = handle(event=event, dense_buffer=...)
```

This removes the current two-step behavior where local materialization happens
before `launch_pending_reduce()` and only the reduce is moved to the prefetch
stream.  If we keep `launch_pending_reduce()` as a helper, it should assume the
local materialization is already enqueued on the same prefetch stream, not on the
forward stream.

### 7.5 First-layer behavior

Layer 0 / first layer is special:

- There is usually little or no preceding layer compute to hide its prefetch.
- It may still need to wait for L2->L1 load before attention can consume cache-hit
  prefix.

The design should still route first-layer initialization through the same
prefetch-safe API so the behavior is uniform and does not accidentally attach
load waits to the wrong stream.  We should not promise speedup for layer 0; the
main objective is avoiding unnecessary serialization for layer 1+.

### 7.6 Fallback and failure policy

Expected no-op cases may stay quiet or debug-level:

- no prefetcher because CP shared KV is disabled,
- prefix too small for prefetch threshold,
- layer is outside this pool.

Unexpected cases must be warning-level and explicit:

- prefetch-safe getter unavailable while HiCache layer transfer is active,
- async reduce unavailable,
- any exception in local prefix materialize/reduce launch,
- falling back from prefetched prefix to synchronous full materialize for layer >
  start_layer.

The warning should include enough fields to distinguish L2->L1 conflict from
normal prefix miss:

```text
cp_rank, layer_id, start_layer, has_layer_transfer_counter,
consumer_indices, prefix_pages, total_slots, path=mla|index
```

Do not add environment gates for correctness warnings.  Temporary high-frequency
profiling logs can be gated or removed after use.

### 7.7 Tests before CUDA ETE

Local/unit tests should not require CUDA kernels.  The important behavior can be
locked with fakes:

1. **No synchronizing getter in prefetch path**
   - fake pool where `get_key_buffer()` raises,
   - `get_key_buffer_for_prefetch()` records the stream argument and returns a
     fake tensor,
   - verify `CpSharedKVMlaPrefetcher.start_next_layer_prefix()` does not call the
     blocking getter.

2. **Layer event is attached to prefetch stream**
   - fake `LayerDoneCounter.wait_until_on_stream(layer, stream)` records calls,
   - verify `next_layer_id - start_layer` is used.

3. **Index path parity**
   - same test for `get_index_k_with_scale_buffer_for_prefetch()`.

4. **Consume path remains blocking-safe**
   - normal `consume()` still accepts the already-prefetched handle,
   - sync fallback still uses normal getter/materialize and emits warning when it
     happens after start layer.

Remote CUDA/ETE validation should then check:

- no warnings about prefetch-safe getter fallback,
- reduced or eliminated `consume_miss` for layer > start layer on cache-hit runs,
- no new collective beyond the existing prefix reduce collectives,
- no worse MTP accept rate after cache hit,
- first-layer wait remains expected, later-layer waits shrink.

### 7.8 Non-goals for this change

Do not solve these in the first patch:

- true `batch_size > 1` CP shared-KV prefill,
- changing radix tree or host-cache residency policy,
- changing owner-lane allocation semantics,
- making current/suffix KV prefetchable,
- adding new CP/global synchronization to prove readiness.

This patch should only fix the composition between existing HiCache layer-load
readiness and existing shared-KV prefix prefetch.

---

## 8. L1 capacity reservation plan for L2->L1 load prefetch

### 8.1 Current code path

Current cache-hit load-back path is:

```text
schedule_policy.py
  AddReq.add_one_req(...)
    -> tree_cache.init_load_back(
         InitLoadBackParams(last_host_node, host_hit_length, mem_quota)
       )

hiradix_cache.py
  HiRadixCache.init_load_back(...)
    -> load_back(last_host_node, mem_quota)

  HiRadixCache.load_back(...)
    -> collect evicted host-backed nodes_to_load
    -> inc_lock_ref(ancestor_node) for quota/protection during planning
    -> cache_controller.load_cp(nodes_to_load)
    -> on owner-lane allocation failure, retry after _evict_for_compute_owner_lanes(...)
    -> set loaded_node.value to reserved device indices
    -> ongoing_load_back[last_hit_node.id] = last_hit_node
    -> inc_lock_ref(last_hit_node) until load ACK

scheduler.py
  new_batch.hicache_consumer_index = tree_cache.ready_to_load_host_cache()

hiradix_cache.py
  ready_to_load_host_cache()
    -> cache_controller.start_loading()

cache_controller.py
  start_loading()
    -> submit per-layer host->device copies on load_stream
    -> record LayerDoneCounter event for each target layer
    -> append ack_load_queue

hiradix_cache.py
  loading_check()
    -> query finish_event
    -> dec_lock_ref(loaded node)
```

The important existing properties are good:

- Device pages are allocated before async H2D starts.
- `loaded_node.value` is assigned before forward, so radix state knows the node is
  now device-resident.
- `last_hit_node` is locked while H2D is in flight, so normal device eviction
  should not evict the just-loaded prefix before the request consumes it.
- CP load uses `alloc_pages_with_owners(page_owners)` to reproduce the write-time
  owner-lane pattern.

### 8.2 Current gap

The current planning is still mostly total-token based:

```text
schedule_policy._get_load_back_mem_quota(...)
  available_device_tokens + evictable_tokens - request_reserve_tokens
```

For CP shared KV this is insufficient.  Load-back capacity is owner-lane
constrained:

```text
required_by_owner[o] <= available_by_owner[o]
```

not just:

```text
sum(required_by_owner) <= sum(available_by_owner)
```

Current `load_cp()` handles this reactively:

1. try `alloc_pages_with_owners(page_owners)`,
2. if it returns `None`, call `_evict_for_compute_owner_lanes(...)`,
3. retry allocation.

That retry path is better than plain total-token eviction, but it is still not a
complete capacity plan:

- the first allocation attempt is expected to fail under lane skew;
- `_evict_for_compute_owner_lanes()` computes lane deficits, but currently calls
  normal `tree_cache.evict(EvictParams(num_tokens=...))`, so eviction candidates
  are still selected by the cache eviction policy rather than by contribution to
  the deficient owner lanes;
- it may over-evict healthy lanes and still fail to free enough pages in the
  deficient lane;
- failure diagnostics do not consistently show `required/available/deficit` per
  owner lane.

The explicit design decision is to replace this scalar admission condition:

```text
available_tokens + evictable_tokens >= required_tokens
```

with a CP owner-lane vector condition:

```text
required_pages_by_owner:  int[cp_size]
available_pages_by_owner: int[cp_size]
deficit_pages_by_owner:   int[cp_size]

all(deficit_pages_by_owner[o] == 0 for o in range(cp_size))
```

This is a per-owner-lane model, not a new cross-rank synchronization model.  All
CP ranks should derive the same `page_owners` and the same logical allocator page
ownership.  If ranks disagree, that is allocator/radix state divergence and must
be surfaced as a correctness bug rather than hidden behind another collective.

### 8.3 Desired invariant

L2->L1 prefetch should never make capacity decisions per layer.

Before submitting any H2D copy, the scheduler/cache layer must have completed:

```text
host-hit nodes_to_load selected
  -> page_owners collected
  -> required_pages_by_owner / available_pages_by_owner / deficit_pages_by_owner computed
  -> owner-lane vector capacity checked
  -> owner-lane deficits evicted, if possible
  -> device pages reserved with matching owner pattern
  -> radix nodes marked device-resident and locked
  -> only then submit async per-layer load
```

Per-layer load and later shared-KV prefix prefetch should only consume already
reserved device slots and local CUDA readiness events.

### 8.4 Proposed implementation units

#### Unit A: owner-lane load-back planning helper

Add a small helper around CP load-back planning, probably in
`hiradix_cache.py` or `common.py`:

```python
@dataclass
class CpLoadBackPlan:
    page_owners: list[int]
    required_by_owner: list[int]
    available_by_owner: list[int]
    deficit_by_owner: list[int]
    host_hit_len: int
```

Responsibilities:

- collect `page_owners` from `node.cp_hicache.page_owners`,
- compute `required_by_owner`, `available_by_owner`, `deficit_by_owner` using
  allocator APIs,
- validate that every node has CP HiCache metadata,
- produce warning fields for all fail-closed cases.

This helper is read-only and should not allocate or evict.

#### Unit B: owner-lane-aware device eviction

Replace the retry-only `_evict_for_compute_owner_lanes()` behavior for CP
HiCache load-back with a real owner-lane-aware eviction path.

The owner-lane-aware eviction should:

1. inspect evictable radix leaves whose `lock_ref == 0`,
2. for each candidate with `node.value`, compute the owner-lane contribution of
   its pages:

```text
page_id = logical_token_loc // page_size
owner = (page_id - 1) % cp_size
```

3. rank candidates primarily by contribution to deficient lanes and secondarily
   by existing eviction policy priority,
4. evict candidates until all lane deficits are covered or no useful candidates
   remain,
5. preserve existing write-through/write-back semantics:
   - backed node: `_evict_backuped(node)`,
   - unbacked node under write-through: `_evict_regular(node)`,
   - write-back policy: use existing `write_backup(..., write_back=True)` path
     only with explicit warnings if it becomes synchronous or expensive.

The first patch can be conservative:

- only specialize CP shared-KV + HiCache load-back;
- keep generic eviction unchanged;
- if owner-aware eviction cannot prove progress, fail closed and log a warning.

No collective is needed.  All owner-lane accounting is local because allocator
page ids and CP size are local and deterministic.

#### Unit C: reserve/lock API boundary

Keep `cache_controller.load_cp()` focused on allocation/enqueue bookkeeping, but
make its contract explicit:

```text
load_cp(nodes_to_load) either returns fully reserved device_indices or None.
It must not return partial reservation.
```

`HiRadixCache.load_back()` remains the owner of radix-state transitions:

- temporary planning lock on `ancestor_node`,
- assign `loaded_node.value` after reservation succeeds,
- add `ongoing_load_back[last_hit_node.id]`,
- hold `inc_lock_ref(last_hit_node)` until `loading_check()` sees the H2D finish
  event.

If any exception occurs after reservation but before `ongoing_load_back` is
registered, free the reserved device indices and roll back node state.  This is
important because the current path crosses Python code and async copy submission;
reservation leaks here become GPU KV leaks.

#### Unit D: diagnostics and fail-fast policy

Capacity failure should be warning-level and structured.  Required fields:

```text
node_id
host_hit_len
nodes_to_load
required_by_owner
available_by_owner
deficit_by_owner
evictable_size
protected_size
ongoing_load_back_count
allocator_available_size
allocator_free_pages
allocator_release_pages
```

Expected skip cases may remain info-level:

- host hit below load-back threshold,
- total host hit exceeds configured quota.

Unexpected cases must be warning-level:

- owner-lane reservation failed,
- owner-aware eviction made no progress,
- retry allocation still failed,
- CP metadata missing from a node selected for CP load-back.

### 8.5 Interaction with shared-KV L1 prefix prefetch

This L1 capacity plan is a prerequisite for the Stage-B shared-KV prefix
prefetch work in Section 7.

Once `load_cp()` returns successfully, Stage-B prefetch must assume:

```text
persistent L1 pages exist and are reserved
the reservation satisfied the per-CP-owner-lane vector check
radix nodes are locked against device eviction
only per-layer H2D readiness remains pending
```

Therefore Stage-B must not trigger device eviction.  If Stage-B cannot access a
layer's persistent shard, that is a bug in reservation/event ordering, not a
signal to evict more memory.

### 8.6 Implementation plan

#### P0: Convert CP HiCache load-back admission from scalar to owner-lane vector

Files:

- `python/sglang/srt/managers/schedule_policy.py`
- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `python/sglang/srt/mem_cache/allocator.py`

Change:

- keep generic scheduler budget as a coarse scalar admission check;
- for CP HiCache load-back, treat scalar `mem_quota` only as an upper bound;
- make final load-back admission depend on:

```text
required_pages_by_owner[o]
available_pages_by_owner[o]
deficit_pages_by_owner[o]
```

Acceptance:

- a request is not admitted for CP load-back unless every owner lane can be
  satisfied after targeted eviction;
- scalar total-token availability is never used as the final CP load-back proof;
- no new all-reduce or cross-rank capacity polling is introduced.

#### P1: Document and test current reservation lifecycle

Files:

- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `python/sglang/srt/managers/cache_controller.py`
- unit test under `test/registered/unit/mem_cache/`

Tests:

- CP load-back returns all-or-nothing reservation.
- `loaded_node.value` is assigned only after reservation succeeds.
- `ongoing_load_back` lock is released only through `loading_check()` ack.

#### P2: Add owner-lane load-back stats helper

Files:

- likely `python/sglang/srt/mem_cache/hiradix_cache.py`
- maybe shared helper in `python/sglang/srt/mem_cache/common.py`

Tests:

- page owner list `[0, 1, 1, 3]` produces required counts correctly.
- allocator free/release page distribution produces available counts correctly.
- missing `cp_hicache` metadata fails closed.

#### P3: Add owner-lane-aware eviction for CP HiCache load-back

Files:

- `python/sglang/srt/mem_cache/hiradix_cache.py`
- possibly `python/sglang/srt/mem_cache/common.py` if reusing from extend alloc

Algorithm:

```text
while deficits remain:
  choose evictable leaf with max contribution to positive deficits
  evict it using existing node eviction primitive
  recompute available/deficit
  stop on no progress
```

Tests:

- skewed allocator state where total pages are enough but owner lane 0 is short.
- normal total-token eviction would not guarantee success; owner-aware eviction
  selects a node containing lane-0 pages.
- locked nodes are skipped.

#### P4: Replace reactive retry with preflight plan + targeted eviction

Current shape:

```text
try alloc_pages_with_owners
if None: evict_for_compute_owner_lanes and retry
```

Target shape:

```text
build load-back plan
if deficit: owner-aware evict before allocation
alloc_pages_with_owners once
if None: warning + fail closed
```

A final retry can remain temporarily for safety, but if it triggers it must emit a
warning that the preflight plan was wrong.

#### P5: Integrate with Stage-B prefetch API work

After owner-lane vector reservation is deterministic, implement Section 7:

- prefetch-safe pool getters,
- stream-specific layer-load wait,
- shared-KV prefix materialization on prefetch stream,
- no forward-stream wait just to launch prefetch.

The Stage-B prefetcher should rely on the P0-P4 guarantee:

```text
for every prefetched layer:
  persistent L1 page slots have already been owner-lane reserved
  radix nodes are locked
  only the local per-layer load event may still be pending
```

#### P6: Remote validation

Remote ETE checks should include:

- cache-hit request under owner-lane skew does not OOM if enough evictable pages
  exist in the right lanes;
- warning logs show per-lane fields on failure;
- no new CP/global collective is introduced;
- L2->L1 load ack releases lock;
- shared-KV prefetch no longer calls blocking getters from forward stream;
- MTP accept rate does not drop after cache hit.

### 8.7 Non-goals

- Do not make true batch-size > 1 prefill part of this patch.
- Do not change host-cache eviction policy.
- Do not introduce all-reduce to agree on capacity.
- Do not allow partial load-back of a host-hit node chain.
- Do not let Stage-B shared-KV prefetch perform device eviction.

---

## 9. Implementation status after P5-P8

Implemented on `cjy-cp-refactor`:

- CP HiCache load-back now builds a per-owner-lane plan from persisted
  `node.cp_hicache.page_owners` before reserving L1 pages.
- Load-back owner-lane deficits are handled by targeted device eviction rather
  than by scalar total-token eviction. If the lane vector still cannot be
  satisfied, the path fails closed with `[CP_HICACHE_FALLBACK]` warning fields.
- Cold/extend allocation now tries compute-owner allocation before scalar
  eviction, passes `EvictParams.owner_lane_deficits` into tree eviction, and
  raises recoverable `KVCapacityWaitError` instead of silently falling back to
  legacy allocation when owner lanes remain exhausted.
- Scheduler prefill allocation catches this recoverable wait, releases the
  `PrefillAdder` radix locks, restores the waiting queue/chunked request state,
  and retries later instead of crashing or leaking locks.
- `SGLANG_DISABLE_TAI_OWNER_SELECT` is intentionally not supported. There is no
  env-gated legacy owner-selection fallback in this path; unexpected owner-lane
  exhaustion is surfaced as capacity wait or warning-level fail-closed behavior.
- HiCache layer-load readiness can now be attached to an explicit CUDA stream by
  `LayerDoneCounter.wait_until_on_stream(...)`.
- KV pools expose prefetch-only getters such as
  `get_key_buffer_for_prefetch(layer_id, stream)` and
  `get_index_k_with_scale_buffer_for_prefetch(layer_id, stream)`. Normal getters
  still block on the current stream and remain the attention-consumer-safe API.
- MLA and NSA-index CP shared-KV prefetch now use the prefetch-safe getters,
  enqueue local prefix materialization on the prefetch stream, and enqueue the
  async CP reduce on the same FIFO prefetch stream. This avoids attaching
  HiCache L2->L1 load waits or prefix materialization to the forward stream just
  to launch next-layer prefetch.

Validation performed:

```text
PYTHONPATH=python python -m py_compile \
  python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py \
  python/sglang/srt/managers/cache_controller.py \
  python/sglang/srt/managers/scheduler.py \
  python/sglang/srt/mem_cache/base_prefix_cache.py \
  python/sglang/srt/mem_cache/common.py \
  python/sglang/srt/mem_cache/hiradix_cache.py \
  python/sglang/srt/mem_cache/memory_pool.py

Remote g0034 container:
PYTHONPATH=python python -m pytest \
  test/registered/unit/mem_cache/test_alloc_pages_with_owners.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
  test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q
# 95 passed
```

Remaining risks:

- The new prefetch stream ordering is unit-tested with fake streams and CPU
  tensors. End-to-end CUDA traces are still required to quantify whether the
  forward-stream stall moved from L2->L1 load to prefix materialize/reduce or was
  actually hidden.
- True `ForwardBatch(batch_size > 1)` remains unsupported for CP shared-KV. Use
  multi-slot overlap as the safer next throughput direction.
- Generic non-owner paths still exist for unsupported CP-shared-KV shapes such as
  multi-batch or non-page-aligned short paths. Those are explicit fallback paths,
  not the intended fast path.
