diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md new file mode 100644 index 000000000..77cdfb59d --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_draft_kv_strong_sync_plan.md @@ -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. diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 1dde726eb..1e2d6e220 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index eaf23cbbc..3792502b7 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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: diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 7446038b7..ed46c88f2 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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)) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index ac9a8b0a0..21d0b4a74 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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