Repository navigation
[KV Offload] Emit self-describing events for recurrent state groups - #62613
Conversation
|
GPU check (Qwen3.5-4B,
Two CPU events per SSM group are the two state checkpoints the request left (prompt end and the aligned boundary); the attention group publishes every block. Probe script: |
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
Documentation preview: https://vllm--62613.org.readthedocs.build/en/62613/ |
mergify
Bot
added
documentation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…e presence The KV router dropped every event whose kv_cache_spec_kind was not main attention, so for hybrid models (Mamba / GDN / linear attention) it had no record of where a worker's recurrent state exists (#16794). - zmq_wire: keep `mamba` events (explicit kind, or kind learned per group_idx from an earlier store) and tag them on the Placement / RouterEvent (`kv_cache_group`: optional, serde-default, named MessagePack only, like `state_source`). Windowed attention groups are still dropped. - indexer: new `KvCacheGroupResidency` records per-(worker, group) state presence by positional block hash across tiers, with the matching cut helpers. Tagged events are routed there and never enter the attention radix tree or the worker-local index, so a Mamba removal cannot evict an attention edge. Worker removal and untagged clears forget the presence. Matching is unchanged in this commit; the lookup cut lands separately. Relies on vLLM publishing per-group BlockStored events (the GPU tier does; vllm-project/vllm#62613 adds the offload tier). Signed-off-by: doxu <doxu@nvidia.com>
|
I think this PR is straightforward. It lets Mamba/recurrent groups emit self-describing offload events instead of placeholders. Dynamo consumes them as state-checkpoint residency, not regular attention blocks. Looks good to me. |
|
This pull request has merge conflicts that must be resolved before it can be |
OffloadingEventsTracker treated every group with a load window as a sliding-window group and emitted a block_size=0 placeholder BlockStored for it. MambaSpec groups have a one-chunk window because they depend on a single state, not because their hashes are meaningless: the offloaded chunk is the recurrent state checkpoint at the chunk's last hash. Consumers that index offload residency (e.g. a KV-aware router) therefore could not see where a hybrid prefix is reusable and overestimated it from the attention group. Describe recurrent groups like full-attention chunks (hashes, tokens, parent, block_size, group_idx) with kv_cache_spec_kind "mamba"; windowed attention groups keep the placeholder. Signed-off-by: doxu <doxu@nvidia.com>
364daf3 to
f3bc5bb
Compare
A recurrent group's offloaded chunk holds one state, valid at the chunk's last hash. Emitting the whole hash range and token span let a consumer treat every hash in a multi-hash chunk as a resumable position, and kept a per-chunk token copy alive for each resident state. The BlockStored now carries that last hash, its parent, and no tokens; BlockRemoved fans out to the same single hash. Attention groups are unchanged. Signed-off-by: doxu <doxu@nvidia.com>
|
/ci run |
|
✅ @luckyq, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #98252 for commit ℹ️ This PR is 137 commits behind upstream |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #98252. |
|
/ci retry |
|
✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/ci/builds/98252 |
|
/ci run |
|
✅ Triggered Buildkite CI #98371 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #98570 for commit |
Purpose
OffloadingEventsTrackeremits ablock_size=0placeholderBlockStoredfor every group that has a load window.MambaSpecgroups get a one-chunk window because they depend on a single state, not because their hashes are meaningless: the offloaded chunk is the recurrent state checkpoint at the chunk's last hash. External consumers that index offload residency (a KV-aware router with a CPU/KVCR tier) therefore cannot see where a hybrid prefix is reusable and overestimate it from the attention group (#62510).Changes
events.py:describes_group()returns true for full-attention andMambaSpecgroups; windowed attention groups keep the placeholder.block_hashesis the chunk's last hash only,parent_block_hashis the hash before it,token_idsis empty,kv_cache_spec_kind="mamba",group_idxandblock_sizeas for attention.BlockRemovedfans out to that same single hash. Consumers align it by hash with the attention group's blocks; announcing the whole hash range would let them treat every hash in a multi-hash chunk as a resumable position. The payload is also small: no per-chunktoken_idscopy is retained for recurrent groups.docs/features/kv_offloading_usage.md: one sentence.Test plan
tests/v1/kv_connector/unit/offloading_connector/test_events.py: 51 passed in a CPU-onlyvllm/vllm-openai:nightlycontainer with thisevents.pydropped in; the 4 tiering-config tests in that file fail there identically without the patch (VllmConfighybrid-KV-manager validation in that environment), so they need the CI environment.mamba_cache_mode=align,CPUOffloadingSpec+self_describing_kv_events, ZMQ subscriber: before, the three SSM groups publish onlyblock_size=0events; after, they publish onekv_cache_spec_kind=mambaevent per chunk carrying the chunk's last hash, at the same hash as the attention group's chunk (table in the first comment; re-run on the checkpoint-only payload pending).