Repository navigation
[HiCache] Add KVCR as a direct external linker with prepared residency - #39946
Closed
ishandhanani wants to merge 13 commits into
Closed
ishandhanani wants to merge 13 commits into
ishandhanani wants to merge 13 commits into
Conversation
A KV hint is advisory routing metadata attached by a trusted orchestrator (a router or a scheduler) after it has picked a worker -- never by an application client. This adds the transport only: the envelope travels unchanged from the entrypoint down to the scheduler's Req, and nothing in the core reads a payload. The envelope is a list of independent actions, each naming its own type and version, so a consumer implementing one action type ignores -- rather than rejects -- an action it does not implement sitting next to it. The payload is therefore untyped here; the component that implements an action owns its schema and validates it. No action types are defined by this change. Wire-compatible with vLLM's KVHintsEnvelope (vllm-project/vllm#53423) so one orchestrator can address either engine with the same envelope. Path: GenerateReqInput (HTTP/Engine) and the gRPC KVHintsEnvelope message -> normalization (single, per-batch-item, and parallel-sample expansion) -> TokenizedGenerateReqInput over msgpack IPC -> Req.
Add --unified-cache-external-linker-backend kvcr. KVCRDirectLinker moves device pool pages straight into KVCR-owned DRAM through NIXL on offload and back into adopted GPU indices on restore; no SGLang host pool exists in this mode. Peer reuse first hydrates the local KVCR tier from a router-hinted source, and lookup() only ever reports pages whose residency claims are held. Interface: UnifiedCacheLinker gains optional prepare_request / preparation_ready / release_request / finish_request hooks with no-op defaults, so Mooncake and UMBP scheduling is unchanged. The wrapper runs the device-only match at enqueue and hands the uncached tail to the backend; the scheduler polls readiness per pass with one attention-group MIN reduction and skips requests still preparing. Layout: one KVCR object per physical pool page, an ordered descriptor list over the page's component/layer spans, one named KVCR subpool per span. Keys are <page-hash>#kvcr-linker-v1#<compatibility-digest>#<pool>; the digest covers checkpoint, KV representation, page geometry, physical spans, draft identity, and shard. Capacity gives every pool the same page count within local_dram_bytes_per_worker divided among local ranks. Ownership: one owner thread performs every KVCR call. Offload results return in FIFO order; a boundary is marked stored only when every object succeeded. Loads complete the layer counter only after all bytes landed, and a failed GPU transfer fails the counter and stops new work rather than recomputing over uncertain slots. Deadlines, stale hints, dead peers, and capacity yield a smaller confirmed prefix or a miss; late completions are drained and their claims released. reset() rebuilds the core so a flush clears local residency; KVCR evictions withdraw EXTERNAL inventory and clear the tree's stored mark. Dense MHA/MLA pools get a direct-linker pool group in the plain-KV strategy, with packed draft layers when draft pools exist. Tests: CPU contract and fault tests run the real linker over CPU tensors with a fake NIXL agent (round trip bytes, hinted-but-absent never hits, sparse boundaries, FIFO offload results, stale attempts, tail realign, deadline with late claims, failed load, reset, inventory removal, backpressure, startup rejections) plus dense pool-group assembly.
Snapshot device indices on the owner thread's private stream instead of its default stream, so a copy never waits behind queued compute kernels. Bound reset/close draining by the core's operation deadline and quarantine what remains. Key the wrapper's stored-page map by event hash so external evictions are matched without a backend import. Count admitted pages from the KV pool only. Drop the adapter's unused blocking call and demotion APIs. Add the KVCR direct linker docs page, a launch script, and an explicit-hint peer reuse check. Give the lock-ref unit fixture the linker attribute the finish hook now reads.
Launches control and linker servers, forces GPU eviction and replays a prompt for the local round trip, and drives explicit-hint peer reuse with no-hint, stale-hint, and dead-peer controls. Records cached tokens, greedy outputs, wall time, and the linker stats lines into a JSON report.
… failure On hosts with InfiniBand devices but no GPUDirect RDMA peer memory, UCX refuses to register CUDA memory (ibv_reg_mr: Bad address) and KVCR fails to start. Wrap the failure with the remedy (UCX_TLS for same-node use or nvidia_peermem) and document it.
…tle before scraping
…linker Runs no external cache, HiCache host tier, and the KVCR linker on one GPU with the same prompts, zero-hit and evict-then-replay workloads, repeated with every sample kept so the spread is visible.
The stats ticker read self._adapter from the owner thread; when core construction outlasted the stats interval the first tick ran before the constructor assigned it and the loop logged a fault. Assign the adapter and reset the stats clock before starting the thread.
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Automations to automatically generate PRs for you. |
github-actions
Bot
added
documentation
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
The direct external linker (
--enable-unified-cache-external-linker) moves KV pages between device pools and an external store with no SGLang host tier, but only Mooncake and UMBP implement it. KVCR (the KV Cache Runner, https://github.com/ai-dynamo/kvcr) is the NIXL-based coordinator behind #36409, whose HiCache adapter stages every offload and restore through a SGLang host pool and keeps the same bytes in two host tiers. This PR adds KVCR as a direct linker:The peer path hydrates the target's own KVCR tier first so residency is confirmed and claimed before a GPU cache hit is admitted. It keeps the target-side DRAM hop; it is not a remote-to-GPU transfer.
Builds on #38891 (KV-hint request transport; its commit is included here until it merges). Requires ai-dynamo/kvcr#39 (typed GPU memory registration). Companion: ai-dynamo/dynamo#14992 (capability publication and hint forwarding).
Modifications
Selection.
--unified-cache-external-linker-backend kvcr, configured through--hicache-storage-backend-extra-configlike the other linkers (local_dram_bytes_per_workeris required; the legacylocal_dram_bytesis refused with a migration message).--enable-hierarchical-cacheand--hicache-storage-backendstay incompatible with linker mode.Interface.
UnifiedCacheLinkergains optional hooks with no-op defaults so Mooncake and UMBP scheduling is unchanged:prepares_requests,prepare_request(LinkerRequestContext, transfers),preparation_ready(handle),release_request,finish_request,take_removed_page_hashes,publishes_external_events.UnifiedCacheLinkerWrapper.prepare_requestruns the request's ordinary device-only match at enqueue and hands the uncached tail to the backend without recursing intolookup;sync_preparationreduces readiness with MIN across the attention group so every rank admits the same requests.Scheduler.
_add_request_to_queue(aggregated and prefill) callsprepare_linker_request;_get_new_batch_prefill_rawskips requests whose preparation is not ready while unrelated requests proceed. Separate from theenable_hicache_storageprefetch path; no HiCache controller is built.Backend (
python/sglang/srt/mem_cache/storage/kvcr/):kvcr_direct_linker.py: owner-thread state machine. A source hint is never a hit:prepare_requestqueries KVCR, fetches useful local or hinted candidates in chunks (acquiring residency claims), andlookupexposes only pages whose claims are held. Claims live until the page is delivered to the GPU or the request is released. Deadlines, stale hints, dead peers, and capacity limits yield a smaller confirmed prefix or a miss; late completions are drained and their claims released. Offload results return in FIFO order and a boundary is published only when every object succeeded. Loads complete the layer counter only after all bytes landed; a failed GPU transfer fails the counter and stops new work instead of recomputing over uncertain slots.reset()(flush) rebuilds the core because KVCR has no clear operation; a core that cannot prove quiescence is quarantined with its DRAM. Bounded in-flight preparation, offload, and abandoned bytes.kvcr_adapter.py: one owner thread performs every KVCR call and polls completions, so the rank also serves peers while idle. No network waits or collectives on that thread.kvcr_layout.py: one KVCR object per physical pool page, an ordered descriptor list over the page's component/layer spans, one named KVCR subpool per span, equal page capacity across pools within the per-rank budget (worker total divided by local ranks; unused tail logged). Keys are<page-hash>#kvcr-linker-v1#<compatibility-digest>#<pool>; the digest covers checkpoint, revision, KV dtype, page geometry, physical spans, draft identity, and shard, excluding runtime pointers and worker or pure DP identity.router_hint.py: parseskv.fetch@1.0from the request envelope, realigns the source endpoint to this rank's attention offset, maps 64-bit router hashes to unsigned, and adapts keys for KVCR.Pool assembler. Dense MHA/MLA pools had no direct-linker pool group;
_PlainKvStrategy.build_direct_linker_pool_groupadds one, with packed draft layers when draft pools exist (focused tests added).Events. Successful offloads publish
EXTERNALBlockStored; KVCR evictions publishBlockRemovedand clear the tree's stored mark so write-through can offload again. L1 eviction does not withdraw external inventory.Docs and example.
docs/docs/advanced_features/kvcr_direct_linker.mdx,examples/runtime/kvcr_direct_linker/(launch script and explicit-hint peer check),test/manual/kvcr_direct_linker/(validation and benchmark drivers).Tests
CPU (
base-a-test-cpu), all passing:test/registered/unit/mem_cache/test_kvcr_direct_linker.py(20): runs the real linker over CPU tensors with a fake NIXL agent. Byte round trip; unprepared and hinted-but-absent pages never hit; sparse trailing-window boundaries; FIFO offload results under out-of-order completion; release and finish drain claims; stale attempts cannot satisfy a new attempt; tail realignment after device growth; device eviction recomputes instead of exposing gaps; deadline yields a miss and late claims release; failed GPU load fails the counter and stops new work; reset clears residency; inventory removal reports event hashes; offload backpressure; key identity survives event-hash conversion; hint parsing; config validation; startup rejections.test_unified_cache_linker.py(46) andtest_linker_pool_assembler.py(13, three new dense cases) pass; the widertest/registered/unit/mem_cachesuite passes (one pre-existingmori-dependent file is not importable here).Real GPU validation
Qwen/Qwen3-0.6B, page size 64, greedy decoding, 16 output tokens, 3000-token prompts,
--max-total-tokens 8192, filler traffic forces GPU eviction between the first request and the replay. Every run compares the restored replay against a no-linker control by output text and servedcached_tokens. Reports and commands are intest/manual/kvcr_direct_linker/kvcr_validate.py.kv.fetchhint, output equals the source controlmiss_no_candidates), no peer waitUCX_TLS=cuda_copy,cuda_ipc,sm,tcpwas required on this host: its mlx5 devices have no GPUDirect peer memory, so UCX refuses to register CUDA memory on them (documented; startup names the remedy).Performance
test/manual/kvcr_direct_linker/kvcr_bench.py, H100 NVL, Qwen/Qwen3-0.6B TP1, page 64,--max-total-tokens 16384, 4 prompts x 3000 tokens, one output token (TTFT), 3 repeats, same GPU, sequential configurations.hitreplays prompts after filler traffic evicted them from the GPU. Medians with min-max spread over 12 samples:Memory: the HiCache arm allocated a 3.77 GB host pool (32832 tokens); the KVCR arm allocated no SGLang host pool and 8 GiB of KVCR DRAM (the "equal total host DRAM" control was not run because the KVCR arm is already far slower).
The KVCR arm restores the right bytes but is slower than recompute for this model and prompt length, and its zero-hit TTFT is both high and highly variable. Linker stats for that arm: 8.4 GB offloaded in 77 tasks, descriptor construction 2.6 s total (~34 ms per 46-page task, on the owner thread), owner-thread preparation latency up to 0.20 s, scheduler-observed admission wait up to 0.64 s.
Isolation run (same arm with
max_inflight_offload_bytesset to 1 so every offload is declined): zero-hit TTFT 45.3 ms (39.1-53.2), preparation latency at most 13 ms, admission wait at most 14 ms, no variance. So the asynchronous preparation path adds roughly 30 ms per 3000-token request on this model (owner-thread round trip plus one extra scheduling pass), and the seconds-scale variance in the full run comes from concurrent offload transfers through the NIXL loopback path.HiCache + KVCR (#36409) and HiCache buffer-only + KVCR could not be run at this revision (see below), so this is a backend comparison against the host tier, not a copy-removal A/B.
Local transfer bandwidth through KVCR/NIXL on this host is the limiting factor: 0.3-0.7 GB/s for GPU<->KVCR DRAM regardless of descriptor granularity, against 3.8-4.5 GB/s host<->host and 24-55 GB/s for a pinned cudaMemcpy (measured in ai-dynamo/kvcr#39). The copy-removal benefit this design targets is therefore not realized on this host; that is a KVCR/NIXL same-agent VRAM path issue, reported there.
Not validated here
Checklist
CI States
Latest PR Test (Base): ❌ Run #35205468759
Latest PR Test (Extra): ❌ Run #35205468332
Latest PR Test (AMD ROCm 10): ❌ Run #35205468827