chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Enforce draft KV strong-sync for CP HiCache hits

Browse Source 
CP HiCache host hits must not advertise target residency unless the draft payload is also valid when EAGLE/MTP draft HiCache is attached. This closes target-only metadata paths by making the CP host-valid predicate and load replay fail fast, resets draft host storage with the target host pool, and records the P1-P3 strong-sync plan state.

The page-index validator is restored for CPU/fake-test tensors only, preserving unit-test coverage for malformed page spans without reintroducing CUDA hot-path host sync.

Constraint: CP shared KV + HiCache + EAGLE/MTP cannot safely demote malformed target/draft metadata to an ordinary cache miss

Rejected: keep permissive fallback for missing draft_host_indices | it can look like a successful cache hit while poisoning speculative acceptance

Rejected: re-enable generic CUDA tensor page validation | it can force host sync in the HiCache transfer hot path

Confidence: high

Scope-risk: moderate

Reversibility: clean

Directive: Do not add silent fallback around CP draft HiCache metadata; unexpected target/draft divergence should fail fast with node/rank context

Tested: remote container targeted tests: 5 passed

Tested: remote container files test_cp_hicache_metadata.py and test_hicache_controller_cp.py: 77 passed

Tested: remote container test_page_index_utils.py: 8 passed

Tested: local git diff --check and py_compile for modified Python files

Not-tested: full CP shared KV + HiCache + EAGLE/MTP ETE

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-27 05:23:31 +08:00
parent 8571fe0cd9
commit 71c4f66968
5 changed files with 899 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

681
docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md Normal file
View File

@@ -0,0 +1,681 @@
# NSA Prefill CP HiCache Draft KV Strong-Sync Plan
Date: 2026-05-27
Branch context: `cjy-cp-refactor` after `8571fe0cd9` (`Share CP HiCache host budget across target and draft KV`).
This document narrows the broader async-state design in
`docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md` to the next
implementation target: make EAGLE/MTP draft KV a strict shadow payload of target
KV under CP shared KV + HiCache. The goal is that draft KV never owns an
independent cache lifecycle. Every request-visible residency transition is
decided by the target radix node, and draft payload validity follows that node.
## Goal
Make this invariant explicit and test-protected:
```text
target HOST_VALID <=> draft HOST_VALID when draft HiCache is attached
target DEVICE_VALID <=> draft DEVICE_VALID when draft uses CP shared KV
host load / host backup / host eviction / radix split / rollback are one logical op
```
The draft pool may have a different per-token byte footprint and therefore a
different host token capacity, but it must not make independent admission,
eviction, or visibility decisions.
## Implementation Status
2026-05-27:
- Phase 1 implemented: CP host-valid predicate now fails fast when host-backed
target metadata is malformed or missing required draft host metadata.
- Phase 2 implemented: `load_cp()` rejects missing draft host metadata before
zero-owned rank handling and frees the replay allocation on failure.
- Phase 3 implemented: `HiRadixCache.reset()` clears target and draft host pools
as one logical reset.
- Adjacent controller invariant guard restored for CPU/fake-test tensors only:
page-shaped host/device indices are validated in unit tests without
reintroducing CUDA tensor truth-value checks or hot-path host sync.
## Fail-Fast / Fallback Policy
For this strong-sync path, silent fallback is a bug. CP shared KV + HiCache +
EAGLE/MTP has too many state dimensions for hidden degradation to be safe:
target-only host hits, missing draft metadata, mismatched owner pattern, or
skipped draft state transfer can all look like a successful cache hit while
poisoning speculative acceptance.
Allowed fallback paths must be narrow and explicit:
1. **Capacity is insufficient.**
- Examples: owner-lane allocation returns `None`, target host allocation is
full before backup, no evictable device pages exist.
- Behavior: return the existing capacity signal, retry eviction when designed,
or fall back to recompute only with an explicit log/counter that says this
was capacity-driven.
2. **Feature is intentionally unsupported.**
- Example: CP HiCache with an unsupported storage backend.
- Behavior: raise during initialization or at the first call site.
3. **Malformed strong-sync metadata or state mismatch.**
- Examples: draft HiCache is attached but `draft_host_indices is None`,
CP node has `host_len > 0` without `cp_hicache`, target/draft state buffer
types disagree under draft CP shared KV, or a page-owner replay invariant
is violated.
- Behavior: raise `RuntimeError`/`ValueError` with node id, lengths, and rank
context. Do not silently demote to cache miss.
In short:
```text
expected resource pressure -> explicit capacity signal + log/counter
unexpected invariant break -> fail fast
unsupported combination -> fail fast
```
## Current Code Facts
### 1. Draft pool is already attached to the target HiCache object
- `python/sglang/srt/managers/scheduler.py`
- `_get_draft_token_to_kv_pool_for_hicache()` discovers the draft worker's
token-to-KV pool.
- `init_cache_with_memory_pool()` passes it into `CacheInitParams`.
- `init_disaggregation()` also calls `tree_cache.attach_draft_kv_pool(...)`
as a compatibility path.
- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `HiRadixCache.__init__()` creates the target host pool and, when CP shared
KV + draft are enabled, attaches a draft host pool.
- `attach_draft_kv_pool()` attaches the draft device/host pools to the same
`HiCacheController`.
- `python/sglang/srt/managers/cache_controller.py`
- `attach_draft_pool()` stores `draft_mem_pool_device` and
`draft_mem_pool_host`, and registers the same `LayerDoneCounter` on the
draft device pool.
This means the right structural place for strong-sync is the target
`HiRadixCache`/`HiCacheController`, not a separate draft radix cache.
### 2. CP backup is mostly atomic today
`HiCacheController._write_cp()` already uses target ownership to drive draft
backup:
1. Compute `owned_positions` and local `physical_device_indices` from target
logical locs.
2. Allocate target host slots.
3. If draft HiCache is attached, allocate draft host slots for the same local
physical count.
4. If draft allocation fails, free target host slots and return failure.
5. Queue target and draft `CacheOperation`s with the same
`physical_device_indices`.
6. Return one `CpHiCacheNodeMetadata` containing both `host_indices` and
`draft_host_indices`.
`start_writing()` merges target and draft queues and appends one `HiCacheAck`,
so transfer completion is already a joint target+draft ack at controller level.
Existing tests cover this partially:
- `test_cp_write_with_draft_pool_backs_target_and_draft_locs`
- `test_cp_write_draft_allocation_failure_rolls_back_target_host`
### 3. CP load is mostly atomic today
`HiCacheController.load_cp()` reconstructs target logical device locs using
`alloc_pages_with_owners(page_owners)`. It then filters to local physical slots
by `owned_positions` and, when draft HiCache is attached, queues draft H2D using
same physical device indices.
The key guard already exists:
```python
if self.has_draft_hicache and node.cp_hicache.draft_host_indices is None:
self.mem_pool_device_allocator.free(device_indices)
raise RuntimeError(...)
```
Existing tests cover the happy path:
- `test_cp_load_with_draft_pool_restores_target_and_draft_locs`
### 4. CP metadata split already preserves draft alignment
`CpHiCacheNodeMetadata.split()` splits `owned_positions`, target
`host_indices`, and optional `draft_host_indices` using the same masks. Existing
tests cover this:
- `test_split_keeps_draft_host_indices_aligned_with_owned_positions`
### 5. Host eviction already frees draft host payload when metadata carries it
`HiCacheController.evict_cp_host()` frees target `host_indices` and, if draft
HiCache is attached and `draft_host_indices` exists, frees draft host indices.
`HiRadixCache._evict_host_for_physical_slots()` clears the target node's
`host_len` and `cp_hicache` after controller host eviction. This is the right
logical ownership model.
### 6. PD transfer already treats draft as extra buffers sharing target pages
Both prefill and decode disaggregation queue setup append draft KV buffers after
target buffers when `draft_token_to_kv_pool` is present. The transfer still uses
target page indices; draft data is an additional payload for those pages, not an
independent page allocation decision.
This matches the desired model: target page indices drive transfer, draft
buffers follow.
## Gaps / Risks Found In Current Code
### Gap A: `_node_backuped()` ignores whether draft metadata exists
Current CP predicate:
```python
return node.host_len > 0 and node.cp_hicache is not None
```
When draft HiCache is attached, a target-only metadata node is still considered
host-backed. That can happen after branch migration, partial feature enablement,
older metadata, or a bug in backup. This must not silently become a cache miss:
it indicates an invariant break in the strong-sync path.
Required behavior: if `cache_controller.has_draft_hicache` is true, host-backed
must also require `node.cp_hicache.draft_host_indices is not None`; otherwise
raise with node/rank context.
### Gap B: `HiRadixCache.reset()` clears target host pool but not draft host pool
Current reset calls:
```python
self.cache_controller.reset()
self.token_to_kv_pool_host.clear()
```
The controller reset clears queues/acks but does not clear the draft host pool.
If reset is used after a draft host pool has allocations, target host allocator
returns to empty while draft host allocator may keep old allocation state.
Required behavior: reset must clear both target and draft host pools under the
same logical operation.
### Gap C: host-visible metadata is assigned before async D2H completion
`HiRadixCache.write_backup()` sets:
```python
node.host_len = len(node.value)
node.cp_hicache = metadata
self.ongoing_write_through[node.id] = node
```
before `writing_check()` observes the transfer ack. The node is locked against
eviction, but `_node_backuped()` can still treat it as host-backed. The current
code likely relies on scheduler/event-loop ordering to avoid loading from an
in-flight write, but that is not encoded in the predicate.
Required decision: either keep current behavior and prove load cannot observe
in-flight host metadata, or add an explicit `host_ready`/pending-state guard so
host hits exclude nodes still in `ongoing_write_through`.
### Gap D: two draft state APIs coexist
`HiCacheController` has both legacy fields:
```python
has_draft
mem_pool_device_draft
mem_pool_host_draft
```
and current CP fields:
```python
draft_mem_pool_device
draft_mem_pool_host
has_draft_hicache
```
The CP path uses `has_draft_hicache`, while old non-CP load/write paths check
`has_draft`. This is survivable but easy to misread and easy to misuse when
adding new lifecycle code.
Required behavior: CP strong-sync code should use one predicate only:
`has_draft_hicache`. Legacy non-CP behavior can remain but should not influence
CP node validity.
### Gap E: tests do not yet encode strong-sync invalid states
Current tests cover successful target+draft backup/load and draft allocation
rollback. Missing tests:
- target-only CP metadata raises when draft HiCache is attached.
- zero-owned target-only metadata raises when draft is attached unless it
carries explicit empty `draft_host_indices`.
- `reset()` clears draft host pool.
- `evict_cp_host()` frees draft indices exactly once in strong-sync mode.
- `load_cp()` frees target device allocation when draft metadata is absent
under draft-attached mode.
- layer-ready completion happens after draft and target per-layer load for the
same layer.
## Implementation Plan
This plan is intentionally conservative: it does not create a draft radix cache
or a draft allocator policy. The target radix node and target allocator metadata
remain the only externally visible cache identity; draft/MTP KV is an attached
payload that follows target decisions.
### Phase 0: Preserve the target-owned identity model
Files:
- Read/modify only if tests expose drift:
- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `python/sglang/srt/managers/cache_controller.py`
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`
Contract to preserve:
```text
radix node identity = target TreeNode
request-visible locs = target logical loc tensor
reload allocation decision = target allocator via page_owners
physical payload movement = target physical indices, with draft as extra payload
```
The host metadata does **not** need to restore the exact old logical page ids.
It must restore the owner pattern saved at backup time:
```python
# write side
page_owners = layout.owner_for_logical_pages(logical_pages)
# load side
device_indices = allocator.alloc_pages_with_owners(page_owners)
```
Do not add a draft-side allocator replay path. If draft needs data movement, use
the physical indices derived from the target allocation.
### Phase 1: Codify the strong-sync host-valid predicate
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Add a helper near `_node_backuped()`:
```python
def _node_has_required_draft_hicache(self, node: TreeNode) -> bool:
if not self._uses_cp_hicache:
return True
controller = getattr(self, "cache_controller", None)
if controller is None or not getattr(controller, "has_draft_hicache", False):
return True
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
return False
if metadata.draft_host_indices is None:
raise RuntimeError(
"CP HiCache invariant violation: draft HiCache is attached but "
f"node_id={getattr(node, 'id', '?')} host_len={getattr(node, 'host_len', '?')} "
"is missing draft_host_indices"
)
return True
```
2. Change `_node_backuped()` for CP:
```python
def _node_backuped(self, node: TreeNode) -> bool:
if self._uses_cp_hicache:
return (
node.host_len > 0
and node.cp_hicache is not None
and self._node_has_required_draft_hicache(node)
)
return node.host_value is not None
```
3. Add tests:
```python
def test_node_backuped_rejects_missing_draft_metadata_when_draft_attached(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = SimpleNamespace(has_draft_hicache=True)
node = TreeNode()
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
with self.assertRaisesRegex(RuntimeError, "missing draft_host_indices"):
cache._node_backuped(node)
```
```python
def test_node_backuped_rejects_cp_host_len_without_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = SimpleNamespace(
cp_shared_kv_layout=SimpleNamespace(cp_rank=1)
)
node = TreeNode()
node.id = 124
node.host_len = 64
node.cp_hicache = None
with self.assertRaisesRegex(RuntimeError, "host_len.*without cp_hicache"):
cache._node_backuped(node)
```
```python
def test_node_backuped_accepts_empty_draft_metadata_for_zero_owned_rank(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = SimpleNamespace(has_draft_hicache=True)
node = TreeNode()
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
draft_host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
self.assertTrue(cache._node_backuped(node))
```
4. Keep the existing no-draft test valid by setting either no
`cache_controller` or `has_draft_hicache=False`.
5. Error message requirements:
- include `node_id`;
- include `host_len`;
- include whether `cache_controller.has_draft_hicache` is true;
- include the current CP rank if available.
### Phase 2: Make `load_cp()` fail closed even on zero-owned ranks
Files:
- Modify: `python/sglang/srt/managers/cache_controller.py`
- Test: `test/registered/unit/managers/test_hicache_controller_cp.py`
Reason found during code re-read: `load_cp()` currently checks
`draft_host_indices is None` only after `owned_positions.numel() > 0`. A
zero-owned rank with target-only metadata would skip the guard and queue a
zero-length target load. Normal `_write_cp()` creates empty `draft_host_indices`
for zero-owned ranks, but controller-level API should still reject malformed
metadata consistently. This is an invariant violation, not a cold-prefill
fallback case.
Steps:
1. Move the draft metadata presence check before the zero-owned `continue`:
```python
for node in nodes_to_load:
node_device_indices = device_indices[offset : offset + node.host_len]
offset += node.host_len
metadata = node.cp_hicache
if self.has_draft_hicache:
draft_host_indices = getattr(metadata, "draft_host_indices", None)
if draft_host_indices is None:
self.mem_pool_device_allocator.free(device_indices)
raise RuntimeError(
"CP HiCache draft KV restore requested but node metadata "
"does not contain draft_host_indices"
)
else:
draft_host_indices = None
owned_positions = metadata.owned_positions.to(device_indices.device)
if owned_positions.numel() == 0:
continue
selected_logical_locs = node_device_indices[owned_positions]
physical_chunks.append(
self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
)
host_chunks.append(metadata.host_indices)
if draft_host_indices is not None:
draft_host_chunks.append(draft_host_indices)
```
2. Add a unit test:
```python
def test_cp_load_zero_owned_rejects_missing_draft_metadata_when_draft_attached(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 128, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=3,
draft_host_pool=FakeHostPool(torch.tensor([], dtype=torch.int64)),
draft_mem_pool_device=FakeDevicePool("draft"),
)
node = TreeNode()
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
with self.assertRaisesRegex(RuntimeError, "draft KV restore requested"):
controller.load_cp([node], node_id=33)
self.assertEqual(allocator.frees[0].tolist(), list(range(64, 128)))
```
### Phase 3: Reset target and draft host pools together
Files:
- Modify: `python/sglang/srt/managers/cache_controller.py`
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Add a controller helper:
```python
def clear_draft_host_pool(self) -> None:
if self.draft_mem_pool_host is not None:
self.draft_mem_pool_host.clear()
```
2. Call it from `HiRadixCache.reset()` immediately after clearing target host
pool:
```python
self.token_to_kv_pool_host.clear()
self.cache_controller.clear_draft_host_pool()
```
3. Add a fake draft host pool test that records `clear()` calls and verifies
reset clears both target and draft host pools.
### Phase 4: Expose target-only CP metadata before cache-hit admission
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Audit all `_node_backuped()` call sites after Phase 1:
- `match_prefix()` host-hit walk.
- `_split_node()` metadata split.
- host eviction candidate handling.
- duplicate-backup avoidance.
2. Desired behavior after Phase 1:
- target-only metadata under draft-attached mode raises immediately.
- `_split_node()` does not split malformed target-only host metadata as if it
were valid.
- host eviction should not silently clear malformed metadata. It should raise
or log an error with node id and rank context before any cleanup path is
used by a deliberate recovery routine.
3. Add at least one `match_prefix()`-level test proving a node with target-only
CP metadata raises before it can be returned as `last_host_node` when draft
HiCache is attached. If constructing the full trie is too costly, add a
focused helper test and keep `match_prefix()` behavior covered by the
existing `_node_backuped()` use.
### Phase 5: Decide in-flight backup visibility before deeper async work
Files:
- Read/modify if needed: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Current risk: `write_backup()` assigns `node.host_len` and `node.cp_hicache`
before D2H completion. The node is lock-protected through
`ongoing_write_through`, and `writing_check()` releases the lock only after the
joint target+draft ack. However `_node_backuped()` currently has no explicit
pending-write check.
Two viable options:
1. **Minimal path:** prove the existing lock/event-loop ordering prevents
in-flight host metadata from being selected for load. Add a unit test or a
comment with a concrete call-chain explanation.
2. **Explicit-ready path:** add a pending host-write guard:
```python
def _node_host_write_ready(self, node: TreeNode) -> bool:
return node.id not in self.ongoing_write_through
```
and require it in `_node_backuped()` for request-visible host hits. Then make
`writing_check()` the only place that transitions the node to ready.
Recommendation: implement the minimal path only if the code proof is solid. If
there is any path where `match_prefix()` can observe a pending write, implement
the explicit-ready path before per-layer backup. Pending write is a valid
lifecycle state, so it may return “not ready”; malformed metadata must still
raise.
### Phase 6: Keep internal transfer split, but expose one logical operation
Files:
- Modify tests: `test/registered/unit/managers/test_hicache_controller_cp.py`
- Read only unless tests fail:
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`
Internal transfer may remain separated because target and draft buffers/pools
are separate. The external contract is still one operation:
```text
one target node id
one target page-index stream
one controller ack
one host-valid predicate
one rollback decision
```
Add controller tests:
1. `evict_cp_host()` frees target and draft host indices together.
2. `_write_cp()` zero-owned rank with draft attached returns empty
`draft_host_indices`, not `None`.
3. `start_loading()` submits draft H2D before completing the corresponding
target layer readiness. For the current implementation, this means the test
should inspect fake host-pool call order around:
```text
draft.load_to_device_per_layer(layer=i)
target.load_to_device_per_layer(layer=i)
producer_event.complete(i)
```
4. PD transfer setup remains target-driven: draft KV buffers are appended, but
no independent draft page-index stream is introduced.
5. Under `SGLANG_CP_DRAFT_SHARED_KV=1`, target/draft NSA state type mismatch is
an error. It should not silently skip draft state buffers because that creates
a target/draft cache-hit mismatch.
### Phase 7: Remote validation
Run targeted tests in the remote container because this repo's runtime deps and
CUDA-specific imports are container-sensitive:
```bash
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && python -m pytest test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q'"
```
Then run one short ETE with CP shared KV + HiCache + EAGLE/MTP:
- Confirm server starts.
- Confirm cache-hit count becomes non-zero after warmup.
- Confirm accept length does not collapse immediately after host cache hit.
- Confirm no `CP HiCache draft KV restore requested but node metadata does not
contain draft_host_indices` appears.
- Confirm no target/draft host pool OOM asymmetry appears.
- Confirm no additional draft-side allocator/page-id decision is introduced in
logs or transfer metadata.
- Confirm every fallback in logs is either a capacity fallback or an explicit
unsupported-feature/error path; no generic silent “use cold path” downgrade for
malformed draft/target metadata.
## Code-Read Conclusions Feeding This Plan
1. The previous CP shared-KV design already handled target owner-lane replay via
`page_owners` and `alloc_pages_with_owners()`. That part should not be
redesigned.
2. The earlier EAGLE/MTP design identified shared target/draft indices as an
implicit assumption. The current plan upgrades that assumption into an
explicit invariant with tests.
3. Current code already routes most target/draft backup/load/evict operations
through one target-owned controller path. The remaining work is to close
malformed-metadata and reset/visibility holes, not to add a separate draft
cache lifecycle.
4. Draft host capacity may be larger than target capacity because draft per-token
footprint can be smaller, but draft residency is still bounded by target node
decisions. Capacity difference must not imply independent draft admission.
## Non-goals For This Step
- Do not implement full TBO interaction yet.
- Do not introduce L3/storage backend support for CP HiCache.
- Do not make backup/load fully per-layer async in this step; this plan only
makes draft KV lifecycle correctness explicit before deeper async work.
- Do not add permanent debug env gates for this step unless a failing ETE cannot
be diagnosed by unit tests.
- Do not introduce a draft radix tree, draft owner-lane allocator replay, or
draft-only cache hit decision.
- Do not introduce broad fallback branches that hide invariant violations as
ordinary cache misses.
## Expected Outcome
After this plan, CP HiCache host-hit behavior should be conservative but
correct for MTP/EAGLE:
- A node is host-backed only if target and required draft metadata are both
present.
- Reset cannot leave stale draft host allocations behind.
- Draft host/device transfers are queued only as part of target logical
operations.
- The next async work can build on a single logical cache state instead of
patching target/draft divergence case by case.

79
python/sglang/srt/managers/cache_controller.py
View File

@@ -40,6 +40,7 @@ from sglang.srt.layers.dp_attention import (
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
from sglang.srt.utils import get_device_module
logger = logging.getLogger(__name__)
@@ -723,6 +724,10 @@ class HiCacheController:
self.prefetch_thread.start()
self.backup_thread.start()
def clear_draft_host_pool(self) -> None:
if self.draft_mem_pool_host is not None:
self.draft_mem_pool_host.clear()
def write(
self,
device_indices: torch.Tensor,
@@ -830,6 +835,30 @@ class HiCacheController:
event.record()
self.ack_load_queue.append(HiCacheAck(event, event, [node_id]))
def _validate_cp_hicache_page_indices(
self,
host_indices: torch.Tensor,
device_indices: torch.Tensor,
) -> None:
"""Validate page-shaped CP HiCache indices without CUDA host sync.
The production invariant is construction-based: HostKVCache alloc/free
preserves page-shaped host spans, and CpSharedKVLayout preserves
per-page contiguity when mapping logical locs to physical locs. The
generic validator uses Tensor truth values and would synchronize CUDA
tensors, so keep the hot path sync-free and validate CPU/fake-test
tensors only.
"""
if not host_indices.is_cuda:
validate_page_aligned_token_indices(
host_indices, self.page_size, "host_indices"
)
if not device_indices.is_cuda:
validate_page_aligned_token_indices(
device_indices, self.page_size, "physical_device_indices"
)
def _write_cp(
self,
device_indices: torch.Tensor,
@@ -909,6 +938,20 @@ class HiCacheController:
required_host_slots=len(physical_device_indices)
)
try:
self._validate_cp_hicache_page_indices(
host_indices, physical_device_indices
)
if draft_host_indices is not None:
self._validate_cp_hicache_page_indices(
draft_host_indices, physical_device_indices
)
except Exception:
self.mem_pool_host.free(host_indices)
if draft_host_indices is not None:
self.draft_mem_pool_host.free(draft_host_indices)
raise
self.write_queue.append(
CacheOperation(host_indices, physical_device_indices, node_id, priority)
)
@@ -1017,6 +1060,19 @@ class HiCacheController:
for node in nodes_to_load:
node_device_indices = device_indices[offset : offset + node.host_len]
offset += node.host_len
if self.has_draft_hicache:
draft_host_indices = getattr(
node.cp_hicache, "draft_host_indices", None
)
if draft_host_indices is None:
self.mem_pool_device_allocator.free(device_indices)
raise RuntimeError(
"CP HiCache draft KV restore requested but node metadata "
"does not contain draft_host_indices"
)
else:
draft_host_indices = None
owned_positions = node.cp_hicache.owned_positions.to(device_indices.device)
if owned_positions.numel() == 0:
continue
@@ -1037,16 +1093,7 @@ class HiCacheController:
self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
)
host_chunks.append(node.cp_hicache.host_indices)
if self.has_draft_hicache:
draft_host_indices = getattr(
node.cp_hicache, "draft_host_indices", None
)
if draft_host_indices is None:
self.mem_pool_device_allocator.free(device_indices)
raise RuntimeError(
"CP HiCache draft KV restore requested but node metadata "
"does not contain draft_host_indices"
)
if draft_host_indices is not None:
draft_host_chunks.append(draft_host_indices)
if not host_chunks:
@@ -1070,6 +1117,18 @@ class HiCacheController:
if self.has_draft_hicache:
draft_host_indices = torch.cat(draft_host_chunks)
try:
self._validate_cp_hicache_page_indices(
host_indices, physical_device_indices
)
if draft_host_indices is not None:
self._validate_cp_hicache_page_indices(
draft_host_indices, physical_device_indices
)
except Exception:
self.mem_pool_device_allocator.free(device_indices)
raise
self.load_queue.append(
CacheOperation(
host_indices,

30
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -883,6 +883,7 @@ class HiRadixCache(RadixCache):
TreeNode.counter = 0
self.cache_controller.reset()
self.token_to_kv_pool_host.clear()
self.cache_controller.clear_draft_host_pool()
# Clear per-request tracking dicts
self.prefetch_loaded_tokens_by_reqid.clear()
self.evictable_host_leaves.clear()
@@ -918,9 +919,36 @@ class HiRadixCache(RadixCache):
logger.warning("Hierarchical cache storage backend is not enabled.")
return False
def _node_has_required_draft_hicache(self, node: TreeNode) -> bool:
if not self._uses_cp_hicache:
return True
controller = getattr(self, "cache_controller", None)
layout = getattr(controller, "cp_shared_kv_layout", None)
cp_rank = getattr(layout, "cp_rank", "?")
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
raise RuntimeError(
"CP HiCache invariant violation: "
f"node_id={getattr(node, 'id', '?')} "
f"host_len={getattr(node, 'host_len', '?')} "
f"cp_rank={cp_rank} has host_len without cp_hicache metadata"
)
if controller is None or not getattr(controller, "has_draft_hicache", False):
return True
if getattr(metadata, "draft_host_indices", None) is None:
raise RuntimeError(
"CP HiCache invariant violation: draft HiCache is attached but "
f"node_id={getattr(node, 'id', '?')} "
f"host_len={getattr(node, 'host_len', '?')} "
f"cp_rank={cp_rank} is missing draft_host_indices"
)
return True
def _node_backuped(self, node: TreeNode) -> bool:
if self._uses_cp_hicache:
return node.host_len > 0 and node.cp_hicache is not None
if node.host_len <= 0:
return False
return self._node_has_required_draft_hicache(node)
return node.host_value is not None
def _node_host_len(self, node: TreeNode) -> int:

28
test/registered/unit/managers/test_hicache_controller_cp.py
View File

@@ -541,6 +541,34 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(host_pool.loads, [])
self.assertEqual(len(controller.ack_load_queue), 1)
def test_cp_load_zero_owned_rejects_missing_draft_metadata_when_draft_attached(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=3,
draft_host_pool=FakeHostPool(torch.tensor([], dtype=torch.int64)),
draft_mem_pool_device=FakeDevicePool("draft"),
)
node = TreeNode()
node.host_len = 4
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(RuntimeError, "draft KV restore requested"):
controller.load_cp([node], node_id=33)
self.assertEqual(allocator.owner_alloc_calls, [[0]])
self.assertEqual(allocator.frees[0].tolist(), [64, 65, 66, 67])
self.assertEqual(controller.load_queue, [])
self.assertEqual(controller.draft_load_queue, [])
def test_cp_load_queues_cpu_host_indices_before_backend_moves(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))

92
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -147,6 +147,42 @@ class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
self.assertEqual(draft_tokens, 80)
self.assertLessEqual(target_tokens * 6 + draft_tokens * 6, 1000)
def test_reset_clears_target_and_draft_host_pools(self):
class ClearablePool:
def __init__(self):
self.clear_calls = 0
def clear(self):
self.clear_calls += 1
class Controller:
def __init__(self):
self.reset_calls = 0
self.clear_draft_calls = 0
def reset(self):
self.reset_calls += 1
def clear_draft_host_pool(self):
self.clear_draft_calls += 1
target_pool = ClearablePool()
controller = Controller()
cache = HiRadixCache.__new__(HiRadixCache)
cache.cache_controller = controller
cache.token_to_kv_pool_host = target_pool
cache.prefetch_loaded_tokens_by_reqid = {}
cache.evictable_host_leaves = set()
cache.pinned_size_ = 1
cache.evictable_leaves = set()
cache._record_all_cleared_event = lambda: None
cache.reset()
self.assertEqual(controller.reset_calls, 1)
self.assertEqual(target_pool.clear_calls, 1)
self.assertEqual(controller.clear_draft_calls, 1)
class TestCpHiCacheNodeMetadata(CustomTestCase):
def test_split_zero_len_moves_all_positions_to_child(self):
@@ -528,6 +564,62 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertTrue(cache._node_backuped(node))
def test_node_backuped_rejects_missing_draft_metadata_when_draft_attached(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
has_draft_hicache=True,
cp_shared_kv_layout=types.SimpleNamespace(cp_rank=2),
)
node = TreeNode()
node.id = 123
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
with self.assertRaisesRegex(
RuntimeError, "node_id=123.*cp_rank=2.*missing draft_host_indices"
):
cache._node_backuped(node)
def test_node_backuped_rejects_cp_host_len_without_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=types.SimpleNamespace(cp_rank=1)
)
node = TreeNode()
node.id = 124
node.host_len = 64
node.cp_hicache = None
with self.assertRaisesRegex(
RuntimeError, "node_id=124.*host_len=64.*cp_rank=1.*without cp_hicache"
):
cache._node_backuped(node)
def test_node_backuped_accepts_empty_draft_metadata_for_zero_owned_rank(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
node = TreeNode()
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
draft_host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
self.assertTrue(cache._node_backuped(node))
def test_single_node_write_lock_updates_device_evictable_leaf_set(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False