Sliding Window Attention (AR video gen) + KV cache lifecycle management

[merceod]

What does this PR do?

Adds windowed autoregressive video generation to M*: instead of denoising a whole clip at once, the video is produced window-by-window, and each finished window streams to the decoder while the next one denoises. Two layers:

Engine (model-agnostic):

• PagedAllocationManager.protect_prefix / release_oldest: free the oldest pages of a live request's KV stream (whole-page front compaction) while a protected prefix (e.g. the text-prompt K/V) survives i.e. the core primitive a sliding-window KV lifecycle needs. A prefix_epoch counter lets caches of prefix-derived views (the dense-attention fast path) detect mid-request context changes, and CPU offload/reload round-trips the new state.
• mstar/engine/windowing.py: WindowSchedule (pure window arithmetic; commit spans tile the sequence exactly, schedules pad to whole windows) and WindowedKVSession (drives protect/release against the cache handle so models never hand-roll page bookkeeping).

Cosmos3 (first adopter, opt-in):

• A video_gen_ar walk: the existing denoise loop run window-by-window, emitting each finished window's latents on a streaming edge to a new vae_decoder_ar node in its own partition. The decoder consumes one window per chunk, decodes it behind a re-decoded left context (seam-free with the causal Wan VAE), and assembles the final mp4 so decode of window k overlaps denoising of window k+1.
• Chained conditioning mode (window_mode: "chained"): each window is a full bidirectional denoise conditioned on the previous window's tail via the existing v2v clean-frame machinery, so it works with the released checkpoints and lifts the clip-length ceiling (e.g. 381-frame videos).
• Everything is gated by enable_windowed_video (default off): node set, walks, partitions, and outputs of existing deployments are unchanged, and requests opt in per call (window_frames / overlap_frames via extra_body). New serving config: configs/cosmos3_nano_ar.yaml.

Cross-window attention through committed K/V with eviction (built on the release primitives above) lands next on this branch.

How was it tested?

• Unit tests: test/modular/test_kv_release.py + test_windowing.py (allocator invariants under randomized grow/release sequences, protection rules, offload round-trip, thread-safety; schedule tiling/overlap/padding properties; page-floor release behavior) and windowed wiring + decoder-assembly tests in the cosmos3 suite - the decoder's context-re-decode-and-trim reconstruction is checked against a whole-stream decode oracle.
• Full CPU suite: the failure set is byte-identical to main's baseline after each commit (no new failures; all new tests pass).
• Zero regression on existing serving: fixed-seed outputs on an unchanged config are byte-identical to the recorded baselines (t2i at 320×192 and 832×480, t2v at 320×192), and the windowed-enabled config reproduces the same t2i bytes.
• End-to-end serving (H200): windowed 57-frame (3-window) and 381-frame (16-window) generations: frame counts exact, every output visually inspected (window boundaries seam-free; gradual long-horizon drift is the known limit of conditioning-only chaining and is what the upcoming KV-context mode addresses), server logs clean.
• To be extended as the branch grows: KV-mode parity against a reference implementation, TP/SP configs, 480p/720p tiers, and the second cluster.

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

[stephen-dwq]

nit: should these test files be moved into ./test/?

[merceod]

These predate this PR and mstar/model/cosmos3/tests/ came in with my original cosmos3 scaffold (#121), and this PR only extends the existing files there (pipeline.py is really the fused reference implementation the parity tests use as source of truth rather than a test module).

I will move the whole in-package dir under test/ in a separate PR since we want to standardize.

[stephen-dwq]

I think we can do some reshuffling here:

"ar_tail" is obtained from stream_tail[:,:,-st["ar_ctx"]:] (line 2271), so tail.shape[2] <= ar_ctz. Then, ctx = tail[:, :, -st["ar_ctx"]:] (line 2263) is a noop.

Because it's a noop, torch.cat([ctx, new.to(ctx.dtype)], dim=2) (line 2264) is the same computation as torch.cat([tail, new.to(tail.dtype)], dim=2) (line 2270).

We can hoist the conditional assignment on line 2270 into the if-else statements above.

[merceod]

Good catch! The invariant holds (ar_tail is always stored pre-capped at ar_ctx, so the re-slice was a no-op and both concats computed the same tensor). I just pushed commit 433b26a where this is done and the tail+window concat is built once and serves as both the decode input and the next tail, which also drops a redundant per-window concat. The decoder unit tests pin bit-equality of chunked vs whole-clip decode and pass unchanged, and the served streaming chunks byte-match the previous build on the changed path.

[Gaurav-Shah05]

I went through the diff and I have added my review comments where necessary, but the PR looks good overall. The engine/model split is right because the allocator knows nothing about windows and the model never touches pages.

I liked the protect/release pair specifically because cross-request prefix caching needs the same thing, that is a block of pages at the head of a stream that survives while everything behind it comes and goes. Once a shared prefix is a protected head instead of ordinary request pages, sharing it across requests is a refcount away. I will build on this. I did have 2 questions:

1. We protect the text prompt, which folds in the attention-sink idea from StreamingLLM. Window 0's committed K/V evicts like any other window though, so nothing visual is anchored past the context horizon. StreamingT2V from CVPR '25 argues drift comes from conditioning on the previous chunk's generated frames, which compounds errors, while a fixed early anchor does not degrade. Their anchor is a trained module injecting position-free CLIP features though, so their gains may not transfer.

   Protecting window 0 along with the text is a one-line change in submodules.py _commit_window; I can try it on your long-generation config and compare prompt adherence and subject consistency over the video. If the gain is small, it can be only text.

2. I checked the Cosmos 3 report: training always denoises the full clip at once, with every frame attending to every frame including future ones, with at most 2 clean conditioning frames, and videos up to 400 frames.

   kv mode differs in two ways. First, each window attends only backward, into a clean history much longer than 2 frames. Second, past 400 frames the frame positions themselves are beyond anything the model saw, like an LLM generating past its trained context length. The first would show as the video degrading smoothly from early on; the second as quality holding and then breaking around frame 400. In your kv outputs, did you happen to notice any one of them?

[Gaurav-Shah05]

Just to make sure I understand the reach of this.

mstar/mstar/engine/kv_store.py

Lines 596 to 617 in 15634ae

	def start_async_retrieve(
	self, request_id: str, label: str, seq_info: SequenceInfo
	):
	seq_len = seq_info.seq_len
	state = self.get_state(request_id, label)
	if state.seq_len >= seq_len:
	return # nothing to do
	first_page = state.seq_len // self.config.page_size
	last_page = (seq_len - 1) // self.config.page_size
	self.alloc(request_id, label, seq_len)
	read_info = []
	for page_pos in range(first_page, last_page + 1):
	token_start = 0 if page_pos > first_page else (state.seq_len % self.config.page_size)
	token_end = self.config.page_size if page_pos != last_page else (
	seq_len % self.config.page_size or self.config.page_size
	)
	local_page_idx = state.page_indices[page_pos]
	remote_page_idx = seq_info.page_indices[page_pos]

start_async_retrieve above pulls a delta by positional page index from the consumer's own seq_len, which assumes the producer's page list only ever grew. After a release the list is compacted, so position P names different tokens than when the consumer last read, and a delta pull would splice wrong pages without raising error/warning; a release also drops the producer's seq_len, so the consumer's nothing-to-do check can pass while the streams have diverged. Is that path reachable for windowed labels today? Either way, could we publish released_tokens in the per-label seq info and have the consumer refuse a delta pull on a mismatch, so this fails loudly instead of silently if placement ever moves a windowed request across workers.

[Gaurav-Shah05]

mstar/mstar/engine/cache_manager.py

Lines 679 to 680 in 15634ae

	to_state.seq_len = from_state.seq_len
	to_state.position_id_start = from_state.position_id_start

KVRequestState now describes a stream with five fields and snapshot_all copies two, which leaves it as the one writer of committed content outside the PR's new rules. It does not bump prefix_epoch the way this line now does, so the dense path's only staleness signal misses that writer; and it drops protected_prefix_tokens, so a snapshot of a protected stream lands unprotected and a release_oldest there would free the text prefix. Bagel is the only caller and it never protects, so both are latent. Could we bump to_state.prefix_epoch in snapshot_all to match this rule, and assert from_state.protected_prefix_tokens == 0 and from_state.released_tokens == 0 there until snapshotting such streams is defined?

[Gaurav-Shah05]

__init__ checks overlap_units against window_units but not context_units against overlap_units, and the two interact: released_end(i) <= window(i+1).start holds exactly when context_units >= overlap_units. Below that, the release after window i frees units inside window i+1's span.

For example:
WindowSchedule(total_units=10, window_units=4, overlap_units=2, context_units=1) gives released_end(0) == 3 with window(1).start == 2, so unit 2 is freed and window 1 conditions on a hole.

Cosmos3 cannot reach it, since kv mode rejects a nonzero overlap_frames and chained mode keeps context_units at 0, but this is the engine-level schedule and the class docs both knobs as usable together. Could we raise a ValueError here when context_units is nonzero and below overlap_units?

[Gaurav-Shah05]

mstar/mstar/engine/cpu_page_pool.py

Lines 90 to 95 in 15634ae

	if cpu_pages is None:
	logger.warning(
	"CPU page pool exhausted: cannot offload %d pages for %s/%s",
	n_pages, request_id, label,
	)
	return

This predates the PR, but the signature change lands on it and windowed requests raise its odds: they hold large streams for a long lifetime, which makes them natural eviction victims.

When the CPU pool cannot fit the request, offload_pages warns and returns without recording anything, and returns no value to check. offload_request then frees the GPU pages, zeroes seq_len, and counts them in its freed total, so the worker logs a successful eviction. The label never lands in offloaded, so reload_request skips it and the request resumes with an empty stream; for a windowed request that means the remaining windows denoise against nothing and the output is silently wrong.

Could we return a bool from here and have offload_request free the GPU pages only when it is True? I can open an issue for this if you prefer that.

[Gaurav-Shah05]

The single-request path guards the overrun before decomposing the counter: _forward_image_gen returns a no-op when step_index >= ar_total_iters. The batched path has no equivalent, and the modulo turns the overrun from benign into the worst case: ar_total_iters is an exact multiple of ar_iters_per_window, so a request dispatched one past its count lands on local == 0, can_batch admits it, and the timestep pick lands on timesteps[0], the noisiest step.

So a finished request re-denoises its emitted latents and steps an exhausted scheduler. The output already left on the streaming edge, so this wastes a step rather than corrupting it. Could we mirror the single-request guard before the modulo?

[Gaurav-Shah05]

These are decoded uint8 frames held on the decode device for every window until the last one lands, and in the shipped configs vae_decoder_ar shares its rank with the DiT, so the retention competes with the model still denoising.

At the served limit that is 64 windows of roughly 30 frames, around 2.5 GB retained plus another 2.5 GB transient at the torch.cat, and since generation-side memory is bounded by context_frames, this becomes what actually caps a non-streamed windowed request. The frames cross the shared-memory edge to the data worker anyway. Could we append pixels.cpu() here and concatenate on CPU, so only the streaming path keeps frames on device?

[Gaurav-Shah05]

A numeric domain_id is bounded below but not above, and it lands in an embedding index. cosmos3_model.py passes the resolved value into params, submodules.py turns it into domain_t, and transformer.py DomainAwareLinear uses it to index nn.Embedding(config.num_embodiment_domains, ...), which is 32 wide.

So domain_id: 999 clears validation and reaches an out-of-range embedding lookup on the GPU, which is a device-side assert rather than a request-scoped error, and that takes the worker's CUDA context with it. The sibling knob is already guarded: cosmos3_model.py bounds raw_action_dim against config.max_action_dim a few lines below the call to this function. Could we pass num_embodiment_domains in here and check 0 <= resolved < num_domains, so a bad id fails at submission the same way a bad raw_action_dim does?

[Gaurav-Shah05]

stream_video decides the response shape here but is not a declared field on VideoGenerationRequest, so the endpoint's streaming mode is invisible in the protocol and this branch and create_videos each fish it out of untyped kwargs. stream on ChatCompletionRequest is declared, and both its readers share the one field. Could we declare stream_video: bool = False the same way and have the adapter setdefault it into kwargs?

[Gaurav-Shah05]

This floors each spatial axis where packing.py build_vision_segment takes math.ceil, so they disagree whenever a latent dimension is odd: at 1280x720 packing gives 920 tokens per unit, this gives 880. The oracle's eviction target then runs low, so the reference attends to context the engine already freed, and the 480p/720p tiers from the test plan would fail parity in a way that reads as an engine bug.

The tested 256x256 grid is even, so nothing fails today. Could we take it from the statics this function already builds, state["cond"]["statics"][0]["num_vision_tokens"] // window_units, as submodules.py _prepare_windowed_prefill does? Same coupling on page_size: it defaults to 128 and the caller never passes it, so it matches KVCacheConfig.page_size by coincidence.

Meant to submit as comment, not blocking. Questions are open, review stands.

[vasilevklart]

Mostly looked at engine files and windowed commit path (didnt go too deep on decoderr etc) Ran the two new test files under test/modular/ locally: 25 passed. The protect/release/epoch design is really clean.. Flagged a memory accounting thing on the offload path inline and a nit + question on the TP/SP configs

[vasilevklart]

release_oldest drops dense_prefix_kv when mutatiing the stream (line 725) but offload_request doesnt. It frees the pages and zeroes seq_len although the dense fast path clones stay on GPU. With windowed kv those clones are the whole stream, not only the text prefix (prefix_len = state.seq_len in _build_dense_gen_plan) And since reload restores the matching prefix_epoch, the clones get reused rather than rebuilt, so nothing ever surfaces it, The offload just reclaims less than the freed page count suggests. At 2 * 36 layers * 8 kv heads * 128 head_dim * bf16 thats =~144 kiB per token, so ~~0.9 GB per branch (~~1.9 GB per request) at the 480p context defaults. While the 320x192 e2e runs would only show ~~140 MB It also means a live windowed request holds its committed context twice (pages and clones) with only the pages visible to the pool's. Could offload_request drop dense_prefix_kv like release_oldest does? The None check at cache_manager.py:1487 already handles the rebuild, one re-gather on reload.

[vasilevklart]

Nit: the last windows commit pass looks like unnecessary work? after the final window commits, no later window ever reads that K V (check_stop ends the loop at that same iteration) and the pages are freed at remove_request Locking the commit_window calls behind window_index + 1 < schedule.num_windows would save one full both branch pass and transient pages per request
Or perhaps it is deliberate so the upcoming cross window mode keeps the full stream committed