Skip to content

Cosmos3-edge PR 2: KV retention policy and partial release for windowed generation - #293

Open
merceod wants to merge 8 commits into
mainfrom
engine/kv-window-lifecycle
Open

merceod wants to merge 8 commits into
mainfrom
engine/kv-window-lifecycle

Conversation

@merceod

@merceod merceod commented Sep 22, 2026 •

Copy link
Copy Markdown
Collaborator

Ports the windowing module and the KV release contract from #198 onto the resource-pool engine (resources/kv/manager.py), so windowed autoregressive video can bound the committed context of a long rollout.

What changed

  • RetentionPolicy(context_budget, protected_prefix) on a CacheStream (the existing stub, filled in) and KVManager.set_retention(request_id, policy, label=None). commit applies the policy right after a stream's span is counted. The oldest unprotected whole pages go back to the arena, stored_len drops by the freed count, and generation bumps so the dense backend re-gathers its prefix. This happens inside the commit, under the lock and before commit_done, so the pre-plan of the next step already sees the compacted stream and no admitted plan addresses freed pages.
  • KVManager.protect_prefix / release_oldest for explicit between-step drivers. Refused under an admitted step or while offloaded.
  • mstar/engine/windowing.py with WindowSchedule (pure window arithmetic, from Sliding Window Attention (AR video gen) + KV cache lifecycle management #198) and WindowedKVSession.bind(prefix), which installs the schedule's context horizon as the stream's retention.

Why the release moved into the commit

#198 released pages between steps on the old allocator. With the pool's speculative pre-planning there is no safe moment between steps anymore (the next step is admitted and planned while the current one runs), so the release became a stream-level policy applied at commit. Streaming ASR or duplex speech models can use the same policy.

Tests

test/modular/test_kv_release.py (compaction, page floor, boundary and tail pages, admitted-step refusal, later plans over the compacted stream, page conservation, commit-time retention with sub-page carry-over, validation, reset) and test/modular/test_windowing.py. ruff and the full test/modular pass. The Cosmos3 kv-mode rollout that drives this end to end is verified on H100 in #297.

@merceod
merceod force-pushed the engine/kv-window-lifecycle branch from 77179c7 to 8bc4628 Compare September 23, 2026 23:14

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant