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

configs/cosmos3_nano_ar.yaml

@@ -0,0 +1,20 @@
+model: "cosmos3"
+# Sequence-length hint for the scheduler. The conductor only asserts its
+# presence; the real per-request capacity is the KV pool below.
+max_seq_len: 8192
+kv_cache:
+  max_num_pages: 1024
+# Windowed-AR serving config: same nano deployment as cosmos3_nano.yaml plus
+# the opt-in windowed video walk (enable_windowed_video adds the
+# vae_decoder_ar node and its streaming decoder partition; requests opt in
+# per call via window_mode). Windowed requests run the eager denoise path, so
+# capture stays enabled only for the t2i tiers it already covers.
+model_kwargs:
+  cuda_graph: true
+  graph_max_latent_area: 2000
+  enable_windowed_video: true
+node_groups:
+  - node_names: ["dit"]
+    ranks: [0]
+  - node_names: ["vae_encoder", "vae_decoder", "vae_decoder_ar", "audio_decoder"]
+    ranks: [0]

configs/cosmos3_nano_ar_sp2.yaml

@@ -0,0 +1,18 @@
+model: "cosmos3"
+# Sequence-length hint for the scheduler (see cosmos3_nano.yaml).
+max_seq_len: 8192
+# Per-rank KV pool (see cosmos3_nano_sp2.yaml).
+kv_cache:
+  max_num_pages: 1024
+# cosmos3_nano_sp2.yaml plus the opt-in windowed video walk: the DiT (and its
+# windowed commit pass) runs Ulysses sequence-parallel across two ranks; the
+# streaming window decoder is small and runs un-sharded on rank 0, consuming
+# the window-latents stream the multi-rank producer dedups to rank 0.
+model_kwargs:
+  enable_windowed_video: true
+node_groups:
+  - node_names: ["dit"]
+    ranks: [0, 1]
+    sp_size: 2
+  - node_names: ["vae_encoder", "vae_decoder", "vae_decoder_ar", "audio_decoder"]
+    ranks: [0]

configs/cosmos3_nano_ar_tp2.yaml

@@ -0,0 +1,18 @@
+model: "cosmos3"
+# Sequence-length hint for the scheduler (see cosmos3_nano.yaml).
+max_seq_len: 8192
+# Per-rank KV pool (see cosmos3_nano_tp2.yaml).
+kv_cache:
+  max_num_pages: 1024
+# cosmos3_nano_tp2.yaml plus the opt-in windowed video walk: the DiT (and its
+# windowed commit pass) runs tensor-parallel across two ranks; the streaming
+# window decoder is small and runs un-sharded on rank 0, consuming the
+# window-latents stream the multi-rank producer dedups to rank 0.
+model_kwargs:
+  enable_windowed_video: true
+node_groups:
+  - node_names: ["dit"]
+    ranks: [0, 1]
+    tp_size: 2
+  - node_names: ["vae_encoder", "vae_decoder", "vae_decoder_ar", "audio_decoder"]
+    ranks: [0]

mstar/api_server/entrypoint.py

@@ -486,6 +486,7 @@ async def iter_result_chunks(self, request_id: str):
         """
         start = time.time()
         finished = False
+        error_delivered = False
         try:
             while True:
                 if time.time() - start > self.timeout_seconds:
@@ -505,6 +506,7 @@ async def iter_result_chunks(self, request_id: str):
                         done = True
 
                 for chunk in new_chunks:
+                    error_delivered |= chunk.modality == "error"
                     yield chunk
 
                 if done:
@@ -522,12 +524,20 @@ async def iter_result_chunks(self, request_id: str):
                     if finished_req is not None:
                         self._finalize_profile(finished_req)
                     for chunk in remaining:
+                        error_delivered |= chunk.modality == "error"
                         yield chunk
                     # A request can fail after the stream is already open
                     # (preprocess error, result-delivery timeout); the HTTP
                     # status is committed by then, so the error must travel
-                    # in-band as the final chunk.
-                    if finished_req is not None and finished_req.error is not None:
+                    # in-band as the final chunk. A preprocess failure is
+                    # already delivered as an in-band error chunk above — only
+                    # synthesize one for errors that never produced a chunk
+                    # (e.g. result-delivery timeout), not a duplicate.
+                    if (
+                        finished_req is not None
+                        and finished_req.error is not None
+                        and not error_delivered
+                    ):
                         yield ResultChunk(
                             request_id=request_id,
                             modality="error",

mstar/api_server/openai/router.py

@@ -129,6 +129,10 @@ async def videos_generations(request: VideoGenerationRequest):
     except Exception as e:  # noqa: BLE001
         default_status = 400 if isinstance(e, (ValueError, TypeError)) else 500
         return _error(getattr(e, "status_code", default_status), str(getattr(e, "detail", e)), "server_error")
+    if (request.model_extra or {}).get("stream_video"):
+        return StreamingResponse(
+            result, media_type="application/x-ndjson", headers={"Cache-Control": "no-cache"}
+        )
     return JSONResponse(result)
 
 

mstar/api_server/openai/serving_videos.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import json
 import logging
 
 from mstar.api_server import media_io
@@ -15,15 +16,18 @@ async def create_videos(api, model_name, adapter, req):  # noqa: ARG001
     args = adapter.video_to_request(req, api.upload_dir)
     request_id = rid("vid")
 
+    stream = bool(args.model_kwargs.get("stream_video"))
     api.submit_request(
         text=args.text,
         file_paths=args.file_paths,
         input_modalities=args.input_modalities,
         output_modalities=["video"],
         model_kwargs=args.model_kwargs,
-        streaming=False,
+        streaming=stream,
         request_id=request_id,
     )
+    if stream:
+        return _stream_ndjson(api, request_id)
 
     chunks = await api.collect_results(request_id)
     # Each video chunk is an mp4 (H.264); return it base64-encoded, mirroring the
@@ -52,3 +56,33 @@ async def create_videos(api, model_name, adapter, req):  # noqa: ARG001
                 logger.exception("Muxing generated audio into the mp4 failed; returning video only")
         data.append({"b64_json": base64.b64encode(video).decode("ascii"), "url": None})
     return {"created": now(), "data": data}
+
+
+async def _stream_ndjson(api, request_id):
+    """Yield one NDJSON line per result chunk for a ``stream_video`` request.
+
+    A windowed request emits each window's mp4 as its own ``video`` chunk;
+    the lines use the ``/generate`` wire shape (modality / base64 data /
+    metadata) with a running ``index``, and the stream is closed by a ``done``
+    line carrying the chunk count. Failures after the stream is open travel
+    in-band as a terminal ``error`` line (the HTTP status is committed), in
+    which case no ``done`` line follows.
+    """
+    index = 0
+    failed = False
+    async for chunk in api.iter_result_chunks(request_id):
+        metadata = dict(chunk.metadata or {})
+        if chunk.modality == "error":
+            failed = True
+        else:
+            metadata["index"] = index
+            index += 1
+        yield json.dumps({
+            "modality": chunk.modality,
+            "data": base64.b64encode(chunk.data).decode("ascii"),
+            "metadata": metadata,
+        }) + "\n"
+    if not failed:
+        yield json.dumps(
+            {"modality": "done", "data": "", "metadata": {"chunks": index}}
+        ) + "\n"

mstar/engine/cache_manager.py

@@ -548,6 +548,10 @@ def advance_seq_len(self, n: int | None = None, pos_id_n: int | None = None) ->
         for rid in self.request_ids:
             state = self._get_state(rid)
             state.seq_len += n
+            if n:
+                # Committed content grew — signal caches of prefix-derived
+                # views (see KVRequestState.prefix_epoch).
+                state.prefix_epoch += 1
             state.position_id_start += (pos_id_n if pos_id_n is not None else n)
 
     @torch.compiler.disable
@@ -601,6 +605,8 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
                     n = seq_lens[i]
                     state = self._get_state(rid, label=label)
                     state.seq_len += n
+                    if n:
+                        state.prefix_epoch += 1
                     if pos_id_ns is None:
                         state.position_id_start += n
                     elif isinstance(pos_id_ns, int):
@@ -614,6 +620,8 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
                 n = ps.seq_lens[i]
                 state = self._get_state(rid, label=label)
                 state.seq_len += n
+                if n:
+                    state.prefix_epoch += 1
                 if pos_id_ns is None:
                     if ps.custom_pos_advance is not None:
                         state.position_id_start += ps.custom_pos_advance[i]
@@ -628,6 +636,26 @@ def advance_seq_lens(self, pos_id_ns: list[int] | int | None = None) -> None:
         for ps in self._plan_states.values():
             ps.custom_pos_advance = None
 
+    @torch.compiler.disable
+    def protect_prefix(
+        self, request_id: str, num_tokens: int, label: str | None = None,
+    ) -> None:
+        """Mark the first ``num_tokens`` of the request's stream (active label
+        unless given) as never releasable. See
+        ``PagedAllocationManager.protect_prefix``."""
+        label = label or self.active_labels.get(request_id, "main")
+        self.alloc_manager.protect_prefix(request_id, label, num_tokens)
+
+    @torch.compiler.disable
+    def release_oldest(
+        self, request_id: str, num_tokens: int, label: str | None = None,
+    ) -> int:
+        """Free the oldest unprotected tokens of a live request, whole pages
+        only; returns tokens actually freed. See
+        ``PagedAllocationManager.release_oldest``."""
+        label = label or self.active_labels.get(request_id, "main")
+        return self.alloc_manager.release_oldest(request_id, label, num_tokens)
+
     @torch.compiler.disable
     def snapshot_all(
         self, from_label: str,
@@ -1460,14 +1488,22 @@ def _run_dense_gen(
         offset = 0
         for idx, prefix_len, gen_len, state in dg["segs"]:
             prefix_cache = state.dense_prefix_kv
-            if prefix_cache is None:
-                prefix_cache = state.dense_prefix_kv = {}
-            cached = prefix_cache.get(layer_idx)
+            if prefix_cache is None or prefix_cache.get("epoch") != state.prefix_epoch:
+                # First gather, or the committed content mutated since the
+                # last one (windowed generation appends/evicts per window —
+                # prefix_epoch moves with every such mutation): drop the stale
+                # clones and re-gather from the live pages. Static-prefix
+                # requests keep the single-gather behavior (their epoch never
+                # moves after prefill).
+                prefix_cache = state.dense_prefix_kv = {
+                    "epoch": state.prefix_epoch, "layers": {},
+                }
+            cached = prefix_cache["layers"].get(layer_idx)
             if cached is None:
                 sub = kv_layer[idx]  # [n_pages, 2, page_size, num_kv_heads, head_dim]
                 k_pref = sub[:, 0].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
                 v_pref = sub[:, 1].reshape(-1, num_kv_heads, head_dim)[:prefix_len].clone()
-                prefix_cache[layer_idx] = (k_pref, v_pref)
+                prefix_cache["layers"][layer_idx] = (k_pref, v_pref)
             else:
                 k_pref, v_pref = cached
             k_parts.append(k_pref)

mstar/engine/cpu_page_pool.py

@@ -16,10 +16,19 @@
 
 @dataclass
 class OffloadedState:
-    """Tracks a single (request, label) that has been offloaded to CPU."""
+    """Tracks a single (request, label) that has been offloaded to CPU.
+
+    Carries the full stream bookkeeping (not just seq_len/position) so a
+    reload restores a windowed label's protection/release state exactly,
+    rather than relying on the live KVRequestState object surviving the
+    offload untouched.
+    """
     cpu_page_indices: list[int]
     seq_len: int
     position_id_start: int
+    protected_prefix_tokens: int = 0
+    released_tokens: int = 0
+    prefix_epoch: int = 0
 
 
 class CPUPagePool:
@@ -68,6 +77,9 @@ def offload_pages(
         gpu_page_indices: list[int],
         seq_len: int,
         position_id_start: int,
+        protected_prefix_tokens: int = 0,
+        released_tokens: int = 0,
+        prefix_epoch: int = 0,
     ) -> None:
         """Copy GPU pages → CPU pages (async on dedicated stream)."""
         n_pages = len(gpu_page_indices)
@@ -94,6 +106,9 @@ def offload_pages(
             cpu_page_indices=cpu_pages,
             seq_len=seq_len,
             position_id_start=position_id_start,
+            protected_prefix_tokens=protected_prefix_tokens,
+            released_tokens=released_tokens,
+            prefix_epoch=prefix_epoch,
         )
 
     def reload_pages(
@@ -102,10 +117,11 @@ def reload_pages(
         label: str,
         gpu_kv_cache: torch.Tensor,
         gpu_page_indices: list[int],
-    ) -> tuple[int, int]:
+    ) -> OffloadedState:
         """Copy CPU pages → GPU pages (async), free CPU pages.
 
-        Returns (seq_len, position_id_start) that were saved during offload.
+        Returns the ``OffloadedState`` saved during offload (its
+        ``cpu_page_indices`` are freed and no longer meaningful).
         """
         state = self.offloaded[request_id][label]
 
@@ -117,11 +133,10 @@ def reload_pages(
                 )
 
         self.page_allocator.free(state.cpu_page_indices)
-        seq_len, pos_id = state.seq_len, state.position_id_start
         del self.offloaded[request_id][label]
         if not self.offloaded[request_id]:
             del self.offloaded[request_id]
-        return seq_len, pos_id
+        return state
 
     def sync(self) -> None:
         """Wait for all pending GPU↔CPU copies to complete."""