chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
47cb9b9d116d1c5c112e894110ec94de80f2db6b
sglang/docs/superpowers/specs/2026-05-07-cp-hicache-host-design.md
T

21 KiB
Raw Blame History

CP Shared KV + HiCache Host Integration Design

Goal: Deliver a first reliable integration between NSA CP shared KV and HiCache host caching for GLM 5 inference.

Architecture: Keep scheduler, radix cache, and req-to-token in the CP-group logical address space. Keep device and host KV pools physical and per-rank sharded. Add a CP-aware HiCache controller adapter plus explicit per-node host metadata so only owned physical shards are transferred while radix state remains logically complete.

Tech Stack: Python, SGLang server runtime, NSA NSATokenToKVPool, CPSharedPagedTokenToKVPoolAllocator, CpSharedKVLayout, HiRadixCache, HiCacheController.

Scope

This design covers the first implementation stage only:

  • Support --enable-nsa-prefill-cp-shared-kv with HiCache host write/load/evict.
  • Keep HiCache storage backends unsupported in this mode.
  • Keep radix tree, scheduler admission, and req-to-token state in full logical KV loc space.
  • Keep each CP rank's device KV pool and host KV pool sharded by physical owner pages.

This design does not cover:

  • HiCache storage backend backup or prefetch.
  • Storage key changes for File, NIXL, Mooncake, EIC, HF3FS, or other backends.
  • Decode-side CP/shared KV.
  • HiSparse NSA KV pools.
  • A full replacement of HiRadixCache with upstream UnifiedRadixTree or HybridCacheController.

If CP shared KV and hicache_storage_backend are enabled together, startup must fail fast with a clear error. Silent fallback is not acceptable because the existing storage path assumes page-hash sequences and contiguous host indices that do not describe CP-owned sparse shards.

Existing Code Facts

The current CP shared KV branch already separates logical and physical KV locs:

  • python/sglang/srt/mem_cache/cp_shared_kv_layout.py provides owner masks and logical-to-physical conversions for torch and NumPy paths.
  • python/sglang/srt/mem_cache/allocator.py::CPSharedPagedTokenToKVPoolAllocator tracks logical allocator size while storing physical pool size separately.
  • python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py profiles physical token capacity and expands logical capacity by CP size.
  • Runtime persistent writes already use CpSharedKVLayout to filter owned logical locs and convert to physical locs before touching local KV buffers.

The current HiCache path is not CP-aware:

  • HiRadixCache.write_backup() passes node.value directly to HiCacheController.write().
  • HiCacheController.write() allocates host indices for the full tensor length and transfers those indices as if they were physical device indices.
  • HiCacheController.load() allocates device indices for the full host tensor length and transfers those indices as if they were physical device indices.
  • HiRadixCache.load_back() uses len(node.host_value) to compute host hit length and to slice restored device indices.
  • HiRadixCache._split_node() slices node.host_value when radix nodes split.

Under CP shared KV, node.value is a full logical loc tensor. It cannot be passed directly into physical host/device transfer functions.

Core Invariants

The implementation must preserve these invariants:

  1. TreeNode.value remains a full logical loc tensor.
  2. req_to_token remains full logical loc state.
  3. Scheduler admission and radix matching remain based on full logical token lengths.
  4. HostKVCache remains physical-only and does not import or depend on CpSharedKVLayout.
  5. CP filtering and logical-to-physical conversion happen in HiCacheController or a small adapter owned by it.
  6. TreeNode.host_value keeps the legacy non-CP meaning: physical host indices tensor.
  7. CP mode does not store CP-owned host subsets in TreeNode.host_value.
  8. CP mode uses explicit metadata to represent host backup state.
  9. Host eviction and load-back accounting report logical token counts, even though local physical transfers are sparse.

Metadata Model

Add explicit CP host metadata for HiRadixCache nodes.

Recommended structure:

@dataclass
class CpHiCacheNodeMetadata:
    logical_len: int
    owned_positions: torch.Tensor
    host_indices: torch.Tensor

    def split(self, split_len: int) -> tuple["CpHiCacheNodeMetadata", "CpHiCacheNodeMetadata"]:
        if split_len < 0 or split_len > self.logical_len:
            raise ValueError(
                f"split_len must be in [0, {self.logical_len}], got {split_len}"
            )
        parent_mask = self.owned_positions < split_len
        child_mask = ~parent_mask
        parent = CpHiCacheNodeMetadata(
            logical_len=split_len,
            owned_positions=self.owned_positions[parent_mask],
            host_indices=self.host_indices[parent_mask],
        )
        child = CpHiCacheNodeMetadata(
            logical_len=self.logical_len - split_len,
            owned_positions=self.owned_positions[child_mask] - split_len,
            host_indices=self.host_indices[child_mask],
        )
        return parent, child

Field semantics:

  • logical_len is the full logical token count represented by this node.
  • owned_positions are node-local logical token positions owned by this CP rank.
  • host_indices are local physical host slots storing the owned token KV data.
  • len(owned_positions) == len(host_indices).
  • owned_positions can be empty on a rank that owns no page in this node.
  • owned_positions are node-local positions, not old logical loc values. This avoids depending on evicted logical locs remaining reusable after load-back.
  • owned_positions and host_indices are stored as CPU torch.int64 tensors. Radix tree operations are CPU-side bookkeeping, and keeping metadata on CPU avoids retaining GPU tensors in long-lived tree nodes.

Add fields to TreeNode or attach equivalent metadata from HiRadixCache:

  • host_len: int = 0
  • cp_hicache: Optional[CpHiCacheNodeMetadata] = None

For non-CP HiCache, existing host_value behavior remains unchanged.

For CP HiCache, backup state is represented by host_len > 0 and cp_hicache is not None.

Radix backup checks must become cache-type aware. A helper is preferable to changing the global TreeNode.backuped property for all cache classes:

def _node_backuped(self, node: TreeNode) -> bool:
    if self._uses_cp_hicache:
        return node.host_len > 0 and node.cp_hicache is not None
    return node.host_value is not None

Use this helper inside HiRadixCache paths that currently read node.backuped. Do not change SWA, Mamba, or plain Radix semantics as part of this integration unless a touched path requires it.

Controller Adapter

HiCacheController is the transfer chokepoint and should own CP-aware translation.

CP HiCache mode must be detected from the allocator or explicit layout, not from the KV pool alone. Use isinstance(token_to_kv_pool_allocator, CPSharedPagedTokenToKVPoolAllocator) or an explicitly passed CpSharedKVLayout to enable the CP adapter. After CP mode is detected, validate that the underlying KV pool is NSATokenToKVPool for this stage.

Add initialization state:

  • uses_cp_hicache: bool
  • cp_shared_kv_layout: Optional[CpSharedKVLayout]
  • page_size: int

The controller can derive CP mode from the allocator by checking for CP-specific fields such as cp_size and cp_rank, or it can receive an explicit layout from HiRadixCache. Prefer explicit construction in HiRadixCache and passing the layout into HiCacheController so the adapter has one source of truth.

CP Write

External call shape stays close to the current API:

metadata = cache_controller.write(device_indices=node.value, node_id=node.id)

CP behavior:

  1. Treat device_indices as full logical locs.
  2. Build owned_mask = layout.owned_by_this_rank(device_indices).
  3. Build owned_positions = owned_mask.nonzero(as_tuple=True)[0].
  4. Build physical_device_indices = layout.logical_locs_to_physical(device_indices[owned_mask]).
  5. Allocate host_indices = mem_pool_host.alloc(len(physical_device_indices)).
  6. If allocation fails, return a failure that exposes required_host_slots=len(physical_device_indices) so HiRadixCache can evict physical host slots and retry.
  7. If the owned count is zero, skip host allocation and skip physical transfer, but still return metadata with logical_len=len(device_indices).
  8. Enqueue or run the existing host transfer using host_indices and physical_device_indices.
  9. Store returned metadata tensors on CPU: owned_positions.cpu() and host_indices.cpu().
  10. Return CpHiCacheNodeMetadata(logical_len, owned_positions, host_indices).

Every logical write-back must produce the same acknowledgement cardinality on every TP rank. If a rank owns zero physical slots for the node, the controller must still create a completed no-op write acknowledgement when HiRadixCache adds the node to ongoing_write_through. This preserves the existing writing_check() invariant that all ranks pop the same logical ack ids after all_reduce(MIN).

The non-CP path continues returning the existing host indices tensor.

CP Load

The load path must allocate full logical device locs but transfer only owned physical shards.

Recommended API:

device_indices = cache_controller.load_cp(nodes_to_load)

CP behavior:

  1. Compute logical_len = sum(node.host_len for node in nodes_to_load).
  2. Allocate device_indices = mem_pool_device_allocator.alloc(logical_len).
  3. For each node, take its slice of device_indices.
  4. Use node.cp_hicache.owned_positions to select the newly allocated logical locs owned by this rank.
  5. Convert selected logical locs to physical locs with layout.logical_locs_to_physical(selected_logical_locs).
  6. Concatenate all owned host indices and owned physical device indices.
  7. If the owned count is zero, skip physical transfer and still return the full logical device_indices.
  8. Enqueue or run the existing host-to-device transfer for the owned physical pairs.
  9. Return the full logical device_indices to HiRadixCache.

CP load-back has the same acknowledgement cardinality requirement as CP write-back. A zero-owned rank must still produce a completed no-op load acknowledgement or equivalent completed producer event so lock release and producer ids remain consistent across ranks.

The non-CP path continues using the existing load(host_indices=host_indices) API shape.

HiRadixCache Changes

Initialization

HiRadixCache.__init__() must detect CP shared KV mode and construct the controller adapter.

It should also enforce:

  • CP shared KV + HiCache storage backend is unsupported and must fail during server argument validation or cache initialization.
  • CP shared KV + HiCache is valid only for NSATokenToKVPool in this stage.
  • CP HiCache mode is enabled from CPSharedPagedTokenToKVPoolAllocator or an explicit CpSharedKVLayout, not merely because the device pool is NSATokenToKVPool.

Backup State

Introduce helpers inside HiRadixCache:

def _node_backuped(self, node: TreeNode) -> bool:
    if self._uses_cp_hicache:
        return node.host_len > 0 and node.cp_hicache is not None
    return node.host_value is not None


def _node_host_len(self, node: TreeNode) -> int:
    if self._uses_cp_hicache:
        return node.host_len
    return len(node.host_value)


def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor:
    if self._uses_cp_hicache:
        return node.cp_hicache.host_indices
    return node.host_value

Use these helpers in paths currently reading node.backuped, len(node.host_value), or node.host_value for host eviction/load-back.

Write Backup

Non-CP path remains unchanged.

CP path:

  1. Call controller CP write with full logical node.value.
  2. If write reports insufficient host slots, evict host entries until physical freed slots are at least the requested required_host_slots, then retry.
  3. On success, set node.host_len = len(node.value) and node.cp_hicache = metadata.
  4. Keep node.host_value is None in CP mode.
  5. Add the node to ongoing_write_through if write-through acknowledgement is required.
  6. Locking and metrics count logical tokens, not owned physical tokens.

_inc_hit_count() must use _node_backuped(node) instead of node.backuped, otherwise CP-backed nodes would appear unbacked because node.host_value remains None in CP mode and the write-through path would repeatedly call write_backup().

Split Node

_split_node() must split CP metadata when a backed-up node is split:

  • Parent metadata keeps positions < split_len.
  • Child metadata keeps positions >= split_len, rebased by subtracting split_len.
  • Parent host_len = split_len.
  • Child host_len = old_host_len - split_len.
  • Host indices are split according to the owned positions mask, not by slicing first split_len entries.

Non-CP host_value slicing remains unchanged.

Load Back

Non-CP path remains unchanged.

CP path:

  1. Collect nodes_to_load while walking evicted ancestors.
  2. Compute host_hit_len = sum(node.host_len for node in nodes_to_load).
  3. Apply load_back_threshold and mem_quota checks to host_hit_len.
  4. Call controller CP load with nodes_to_load.
  5. If allocation fails, evict device logical tokens and retry.
  6. Assign full logical slices from returned device_indices to each node.value using node.host_len.
  7. Update evictable_size_, locks, and metrics by full logical lengths.

The load_back_threshold check intentionally uses logical length in this first stage. Using local physical owned-token count would let different CP ranks make different load-back decisions. A physical-transfer threshold would require a CP-group reduction and is deferred.

Host Eviction

Non-CP path remains unchanged.

CP path:

  1. Host eviction candidate selection still uses radix logical nodes.
  2. Free only node.cp_hicache.host_indices from the local host pool.
  3. Track two counts separately: physical freed slots for allocator retry progress, and logical evicted tokens for metrics/debug output.
  4. Clear node.host_len and node.cp_hicache.
  5. Remove the node from the radix tree as the existing host eviction path does.

Host eviction for CP write retry must progress by physical freed slots, not by logical host_len. A node can have a large logical host_len but zero local host_indices on a rank. Counting that logical length as eviction progress could stop the eviction loop without freeing the physical host slots needed by the retry.

Capacity Policy

Host memory remains per-rank physical capacity based.

The existing HostKVCache default sizing uses device_pool.size * hicache_ratio when hicache_size <= 0. Under CP shared KV, device_pool.size is already physical per-rank size, so this is the desired behavior.

Do not change host pool sizing to logical capacity. That would multiply host memory by CP size per rank and defeat the CP memory-saving goal.

Do not add a new CP-specific host ratio parameter in the first stage. Existing hicache_ratio and hicache_size remain the user-facing controls.

dev-hicache Fix Integration

Do not blindly cherry-pick the dev-hicache commits onto dev-cp-refactor. The CP branch already modified allocator/common/model runner code, and direct cherry-picks risk overwriting CP shared KV logic.

Manually port fixes in this order:

  1. Correctness fixes before CP HiCache integration:
    • Load-back mem_quota scheduling budget from a666ab552.
    • Paged allocator extend OOM accounting from d6ca2fd0c, merged carefully with compute-owner lane eviction in common.py.
    • defaultdict(TreeNode) ghost-node fix from e008d2cf3 for radix-like nodes.
  2. CP host HiCache integration from this design.
  3. HiCache performance fixes from 25b6a957c:
    • Incremental eviction heap.
    • Cached timestamp in hot loops.
    • Persistent all-reduce tensors.
    • Async write acknowledgement thread.

The short bigram page lookup fix from upstream commit 7be96c8ce is a separate small correctness task and should be manually ported with a unit test.

The upstream SWA HiCache storage type registration is out of scope for GLM 5 NSA host-only integration.

Error Handling

Startup or initialization must fail with clear messages for unsupported combinations:

  • CP shared KV + HiCache storage backend.
  • CP shared KV + HiSparse NSA pool.
  • CP shared KV + non-NSA KV pool in HiCache host integration.

Runtime allocation failure handling follows existing HiCache behavior:

  • Host write allocation failure returns None; HiRadixCache evicts host entries and retries.
  • Device load allocation failure returns None; HiRadixCache evicts device entries and retries.
  • If retry still fails, load_back() returns None and the request recomputes the prefix as today.

For CP metadata consistency failures, fail loudly with assertions in unit-testable helper methods. Examples:

  • metadata.logical_len must match the node logical length.
  • len(metadata.owned_positions) == len(metadata.host_indices).
  • split lengths must be within [0, logical_len].
  • owned_positions must be sorted and in range.
  • owned_positions and host_indices must be CPU tensors before being stored on TreeNode.

Testing Strategy

Use SGLang unit-test conventions for new tests:

  • Place tests under test/registered/unit/.
  • Inherit from CustomTestCase.
  • Register CPU tests with register_cpu_ci where possible.
  • Avoid launching a server for metadata and controller mapping tests.

Required unit coverage:

  1. CpHiCacheNodeMetadata.split():
    • page-aligned split.
    • split_len == 0.
    • split with no owned positions on this rank.
    • split where owned positions are on both sides.
    • invalid split values raise errors.
  2. Controller CP write mapping:
    • full logical locs produce only owned host allocation.
    • physical device indices passed to host transfer match CpSharedKVLayout.
    • zero-owned node returns metadata without allocating host slots.
    • zero-owned node still produces one completed logical write acknowledgement if the node enters ongoing_write_through.
  3. Controller CP load mapping:
    • full logical device locs are allocated.
    • owned positions map to newly allocated logical locs, then to physical locs.
    • zero-owned metadata skips transfer but returns full logical locs.
    • zero-owned metadata still produces a completed logical load acknowledgement or producer event.
  4. HiRadixCache metadata behavior:
    • _node_backuped() returns true for CP nodes with host_len and metadata.
    • _inc_hit_count() uses _node_backuped() and does not repeatedly write back already CP-backed nodes.
    • _split_node() splits host_len and CP metadata correctly.
    • load_back() uses host_len rather than len(host_value) in CP mode.
    • host eviction frees owned host indices and uses physical freed slots as retry progress.
  5. Server args/config validation:
    • CP shared KV + hicache_storage_backend fails fast.

Required regression checks:

  • CP shared KV layout tests.
  • Radix cache unit tests.
  • Scheduler prefill-adder tests covering mem_quota.
  • New CP HiCache metadata/controller tests.

Optional manual GPU smoke:

  • Launch GLM5/NSA CP shared KV with HiCache host-only.
  • Use repeated-prefix requests to trigger backup, device eviction, and host load-back.
  • Confirm no logical/physical index assertion failures and no host/storage unsupported path is entered.

Implementation Order

  1. Add tests for existing correctness fixes that are currently missing from the CP branch.
  2. Manually port load-back mem_quota, allocator OOM accounting, and defaultdict(TreeNode) fixes.
  3. Add CP+storage fail-fast validation.
  4. Add CpHiCacheNodeMetadata and unit tests.
  5. Add CP adapter methods in HiCacheController with mock unit tests.
  6. Wire HiRadixCache backup, split, load-back, and host eviction to CP metadata helpers.
  7. Run CPU/unit regression tests.
  8. Manually smoke test GPU CP+HiCache host-only if hardware/model access is available.
  9. Port performance fixes from 25b6a957c as separate patches with regression tests.

Open Follow-Up Work

These items are deliberately deferred:

  • CP-aware HiCache storage key and prefetch design.
  • Support for File/NIXL/Mooncake/EIC/HF3FS storage backends under CP HiCache.
  • Independent CP-specific host ratio configuration.
  • Decode-side CP/shared KV HiCache support.
  • Upstream UnifiedRadixTree or HybridCacheController migration.

Deferred Performance Follow-Up

After host-only CP HiCache correctness lands, port HiCache performance work from 25b6a957c as separate changes with their own tests. Split cached timestamp use, incremental eviction heap maintenance, persistent all-reduce tensors, and async write acknowledgement into independent patches. Each patch must preserve CP zero-owned write/load no-op acknowledgement cardinality and must rerun the CP HiCache unit regression commands from the implementation plan.

Reference in New Issue View Git Blame Copy Permalink