mirror of
https://github.com/kvcache-ai/Mooncake.git
synced 2026-11-05 01:17:55 +00:00
[GH-ISSUE #4803] [RFC]: Per-sequence state objects and two-stage hit resolution for hybrid/SWA models #1216
Labels
No labels
ascend
auto-closed
bug
bug
CI/Build
Common
documentation
enhancement
feature request
good first issue
help wanted
high priority
MetaX
nightly-failure
pull-request
PyTorch Backend
question
RFC
Roadmap
stale
Store
Tests
Transfer Engine
No milestone
No project
No assignees
1 participant
Notifications
Due date
No due date set.
Dependencies
No dependencies set.
Reference
kvcache-ai/Mooncake#1216
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Originally created by @sususama on GitHub (Oct 20, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4803
Changes proposed
Motivation
Mooncake Store models every cache entry as a paged KV block addressed by a content-addressed prefix chain. That covers full-attention models well. It does not cover models where part of the persistent per-sequence cache is not paged and not addressed by block position.
DeepSeek V4.1-Flash is one such model, and Mooncake already touches this family through the Engram work (#4083, #4692). In V4.1 only the first 20 layers keep a sliding-window KV, the window is 128 tokens plus 6 tokens of speculative-decode rollback, and the retained state is a fixed-length per-sequence slot rather than a page. The open-source DeepInfer release documents the layout precisely in
Property | Value -- | -- unit | one fixed-length slot per sequence (StateBlock), not a page persisted content | first 20 layers only persistence granularity | only at 128-token boundaries, and only after the last prefill chunk ring length | 134 (window 128 + rollback 6) on the inference side; 396 on the prefill side bytes per sequence | 20 × 134 × 528 B = 1,415,040 B (MXFP8) medium | host memory of the machine that produced it, plus a directory of which machine holds which keydoc/DeepSeek V4.1 Flash 推理部署 Q&A.pdf:The same shape appears in Gemma3n/Gemma4, Qwen4-Exp and the rest of the hybrid stack that the SGLang adapters in #4692 already cover.
The object model and the hit semantics are the two halves of one problem, which is why they are in one RFC:
Why the current model cannot express it
block_sizetokens of paged KV.Proposed design
A. State object kind
state_layout(per-layer byte offsets).state_layoutis registered per model/checkpoint like a schema, so Store can validate size and reject mismatched producers instead of serving garbage.boundary_aligned_lengthand refuse non-aligned writes, matching the engine-side rule that state is only written on aligned boundaries.B. Two-stage hit resolution
The engine-side counterpart is "replay only the last
windowtokens, with the attention window truncated inside the replay segment", which makes the recompute cost independent of context length. That stays engine work; this RFC is about the storage shape and the reported hit length.C. Lifetime and placement
What this unlocks
Without it, a hybrid/SWA model on Mooncake gets either (a) SWA bytes smuggled into a block object, which inflates payload and breaks block reuse, or (b) no prefix reuse at all for the layers that need it. With it, prefix reuse for these models costs ~1.4 MB per cached sequence and the router stops advertising matches the engine cannot use.
The asymmetry is the reason to bother: a false long hit costs a full prefill, a conservative short hit costs one block of recompute.
Relationship to existing work
Phasing
Steps 1 and 2 are independently reviewable; they are one RFC because step 2 is meaningless without step 1, and the step-back granularity in step 2 depends on the alignment rule chosen in step 1.
Non-goals
Open questions
state_layoutregistry acceptable, or should the value stay fully opaque with size validation only?ceil(window / block_size)blocks? In the reference implementation state is stored at 128-token boundaries with a 128-token window, so one block is enough there; the general rule needs a definition.L1and let the engine decide? (Pro: fewer conservative step-backs. Con: the router no longer knows the real cost.)References
deepinfer/worker/cache_engine.py(StateBlock,select_state_slice_to_store,load_state_from_store_to_prefill_cache) anddoc/DeepSeek V4.1 Flash 推理部署 Q&A.pdf.Before submitting a new issue...
@github-actions[bot] commented on GitHub (Oct 20, 2026):
Thanks for opening this issue, @sususama!
52159774A maintainer will triage this when possible. To help us respond faster, please include:
Useful links: Documentation · Contributing guide
@he-yufeng commented on GitHub (Oct 20, 2026):
The state-validity and routing contract is worth separating from a new Store object kind. Store already accepts arbitrary byte slices under a string key (ObjectKey, Put/Get), so a fixed-size state payload does not by itself require changing that layer. A first adapter could use a separate key namespace containing the checkpoint/layout fingerprint, prefix hash and exact aligned boundary, validate the payload size, and publish availability only after the write completes.
For the lookup result, I would keep
block_hit_lengthseparate from state availability and the engine's fallback cost. An engine that supports bounded replay can recover differently from one that requires exact reconstruction; stepping back one block is not a safe general rule for every hybrid/SWA model. That capability and the replay boundary need to come from the engine rather than be inferred from an object's presence.One source correction for section C: the reported 3/7-day retention is for Block KV. The SWA state in the reference deployment stays in host memory for about 180 seconds. Independent lifetimes are useful here, but those numbers should not become a state-object TTL default. I would pin the adapter's state layout and add hit/missing/expired/mismatched-layout contract tests before deciding whether the Master needs a model-specific layout registry.
@sususama commented on GitHub (Oct 21, 2026):
Thanks, this is a useful correction and I'll fold both points in — and I'll write the adapter.
Section A — agreed, no new object kind.
ObjectKeyis astd::stringand Put/Get already take arbitrary byte payloads, so I'll rewrite this as an adapter contract rather than a Store change: a separate key namespace carrying{checkpoint/layout fingerprint, prefix hash at the aligned boundary}, payload-size validation on both sides, and availability published only after the write completes. Thestate_layoutregistry becomes an adapter-side pinned constant in v1.Section B — agreed, and this is the better shape. Store should report facts, not infer policy. Concretely: return
block_hit_lengthandstate_available(aligned_length)as separate fields and let the engine combine them with its own replay capability. I'll drop "step back one block" as a Store rule. One note on why both halves are needed: an engine that declares bounded-replay support still cannot act on it if the lookup never reports the state as missing.Section C — you're right, I had this wrong. The 3/7-day figures are Block KV retention and should not appear as a state TTL default; I'll remove them. For the record, the open release does not pin a constant here: state capacity is a slot pool sized by
DEEPINFER_NUM_STATE_SLOTS_RATIO, so residency is driven by pool pressure and request lifetime rather than by time. If the ~180 s figure comes from something citable I'd rather reference it than paraphrase it.Contract tests — agreed, and I'd start there rather than argue about the registry: hit / missing / expired / layout-mismatch / boundary-misaligned against a pinned layout.
One question before I open anything: should the adapter live on the
EngramStore-style path, or as a new lightweight connector undermooncake-integration? Default if you would rather not spend time on this: a new lightweight Python module on the store-facing side (python/mooncake/store/), followingstructured_object_store.py, because the adapter needs no GPU row gather andput/get/is_exist/get_sizeare already exposed. I'll proceed that way if I don't hear otherwise; it is easy to move before anything depends on it.@he-yufeng commented on GitHub (Oct 23, 2026):
Confirmed, go with the default: a lightweight module under
python/mooncake/store/followingstructured_object_store.py.mooncake-integrationis the C++/allocator-level surface, and this adapter is exactly the kind of byte-payload contract that belongs on the store-facing Python side. Nothing about it needs the GPU row-gather path.Two things worth pinning before the first PR lands, both cheap now and expensive to retrofit:
state_layoutas an adapter-side constant is right for v1, but put a scheme/version byte (or a short semver string) into the fingerprint so a v2 layout change reads as a clean miss instead of a collision.state/plus the fingerprint) keeps adapter keys out of the way of engine-written Block KV keys, and the contract tests should include a namespace-collision case, not only hit/miss/expired/layout-mismatch/boundary-misaligned.On the TTL point: agreed, leave time-based figures out of the state path entirely. Pool-pressure-driven residency via the slot ratio is the accurate description, and the adapter should treat expiry as an observed fact from
is_exist/get_size, never as a duration it assumes.Looking forward to the contract tests first, as you proposed.
@whn09 commented on GitHub (Oct 23, 2026):
An engine-side consumer for this contract now exists: deepseek-ai/DeepInfer#216 (draft) adds Mooncake Store as an L3 tier under DeepInfer's host prefix cache (#144 there), for DeepSeek-V4.1-Flash. Notes from it that may help the adapter and the contract tests:
How it maps to Sections A/B today (plain
MooncakeDistributedStore, until the adapter lands):state/<fp>/<prefix chain hash at the 128-aligned boundary>. One object per boundary, holding exactly the persisted SWA region DeepInfer stores at that position (state_size_to_disk, taken withselect_state_slice_to_store). Boundaries are the last few 128-aligned positions of each stored prompt.dsv4kv/<fp>/g<G>/<chain hash of the last block>. Each object holds G consecutive 128-token blocks, all layers per block.<fp>:v1-<sha256[:16]>over the checkpointconfig.jsonplus page bytes, state bytes, block size and the env knobs that change the state ring layout. The version is in the fingerprint itself, as suggested, so a layout change is a clean miss.batch_is_existon block keys gives the block run;batch_is_existon state keys at candidate boundaries gives state availability. The engine then picks the deepest boundary that it can also replay from (it must recompute ≥ 140 trailing tokens) and fetches withbatch_get_into_multi_buffers. The Store never infers policy.existsandget(we treat-704/ size mismatch as a miss) falls back to the host-only lookup. There is a unit test for exactly that case.Observations from B200, TCP, one host, wheel 0.3.13.post1 that may matter for defaults:
batch_put_fromreaches 0.17 GB/s with one block per object, but 1.7 GB/s with 4 blocks per object (the loopback limit). Hence G=4 (512 tokens) for block objects. State objects are larger and single.protocol=shm(unsupported_protocol).protocol=efaon the same host gaveTRANSFER_FAIL (-800)for every put; I haven't dug into that yet.Once
python/mooncake/store/'s adapter lands, DeepInfer can switch its state path to it directly. Happy to run the contract tests against real V4.1 states, and to share layout numbers, if that is useful.中文
这个约定现在已经有了引擎侧的使用方:deepseek-ai/DeepInfer#216(Draft)在 DeepInfer 的主机前缀缓存(那边的 #144)下面加了一层 Mooncake Store 作为 L3,面向 DeepSeek-V4.1-Flash。下面几点或许对 adapter 和契约测试有帮助。
目前和 Section A/B 的对应关系(adapter 落地之前,直接用
MooncakeDistributedStore):state/<fp>/<128 对齐边界处的前缀链哈希>。每个边界一个对象,内容正好是 DeepInfer 在该位置保存的 SWA 区域(state_size_to_disk,用select_state_slice_to_store取出)。边界取每个已保存 prompt 的最后几个 128 对齐位置。dsv4kv/<fp>/g<G>/<最后一个 block 的链哈希>。每个对象包含 G 个连续的 128-token block,每个 block 含所有层。<fp>:v1-<sha256 前 16 位>,由 checkpoint 的config.json加上页字节数、state 字节数、block 大小,以及会改变 state 环形布局的环境变量算出。按建议把版本号放进了指纹本身,布局一变就自然查不到。batch_is_exist得到连续 block 的长度,对候选边界的 state key 做batch_is_exist得到 state 是否可用;然后由引擎挑出最深的、自己也能重放的边界(末尾必须重算至少 140 个 token),再用batch_get_into_multi_buffers读取。Store 本身不做任何策略推断。exists和get之间对象消失时(-704或大小不符都按未命中处理),退回只用主机池的查找。这种情况有专门的单元测试。在 B200 上的观察(TCP,单机,wheel 0.3.13.post1),可能和默认参数有关:
batch_put_from只有 0.17 GB/s,每对象 4 个 block 能到 1.7 GB/s(回环带宽上限),所以 block 对象取 G=4(512 token)。state 对象本身较大,单独存放。protocol=shm(unsupported_protocol);同机用protocol=efa时每次 put 都返回TRANSFER_FAIL (-800),我还没来得及深入查。等
python/mooncake/store/的 adapter 落地后,DeepInfer 可以直接把 state 路径切过去。如果有用,我很乐意用真实的 V4.1 state 跑契约测试,也可以提供布局相关的数据。