bug fixes: token order, failed streams, capture order, SHM, configs, prefix cache docs

configs/qwen3omni_2gpu.yaml

@@ -1,5 +1,12 @@
 model: "qwen3_omni"
 max_seq_len: 32768
+# The Thinker's KV pool defaults to 2,048 pages, 24 GiB at 12 MiB a page, which
+# doesn't fit beside its weights on an 80 GB GPU; qwen3omni_2gpu_large_kv.yaml
+# keeps the default for larger GPUs. Host pages give a full pool a request to move.
+resources:
+  thinker_kv:
+    max_num_pages: 1024
+    cpu_offload_pages: 1024
 node_groups:
   - node_names: [audio_encoder, vision_encoder]
     ranks: [1]

configs/qwen3omni_2gpu_large_kv.yaml

@@ -0,0 +1,15 @@
+model: "qwen3_omni"
+max_seq_len: 32768
+# The Thinker's KV pool at its default 2,048 pages, 24 GiB at 12 MiB a page. It fits
+# beside the weights on an H200; qwen3omni_2gpu.yaml sizes it for an 80 GB GPU.
+node_groups:
+  - node_names: [audio_encoder, vision_encoder]
+    ranks: [1]
+  - node_names: [Thinker]
+    ranks: [1]
+
+  - node_names: [Talker]
+    ranks: [0]
+
+  - node_names: [Code2Wav]
+    ranks: [0]

configs/qwen3omni_colocated.yaml

@@ -1,5 +1,12 @@
 model: "qwen3_omni"
 max_seq_len: 32768
+# The KV pools default to 2,048 pages. The Thinker's is 24 GiB at 12 MiB a page,
+# and every node's weights on one 80 GB GPU leave about 12 GiB.
+resources:
+  thinker_kv:
+    max_num_pages: 256
+  talker_kv:
+    max_num_pages: 256
 node_groups:
   - node_names: [audio_encoder, vision_encoder, Code2Wav]
     ranks: [0]

docs/serving.rst

@@ -226,10 +226,13 @@ travels with its other ``model_kwargs``.
 
 A cached prompt no longer reserves its pages, so the pool no longer limits how many
 requests are admitted. It has to hold the decode of every request allowed to run at once,
-and a decode step that finds no free page holds its requests until they time out. Two
-settings prevent that. ``max_concurrent_requests`` caps how many requests run at once, and
-the pool's pages divided by the pages one request needs, prompt and output together, is a
-safe value. ``cpu_offload_pages`` gives a full pool a request to move to the host.
+and a decode step that finds no free page holds its requests until they time out.
+``max_concurrent_requests`` prevents that. It caps how many requests run at once, and the
+pool's pages divided by the pages one request needs, prompt and output together, is a
+safe value. ``cpu_offload_pages`` lets a full pool move a request to the host. An offload
+frees only the device pages no other running request reads, though, so a request that
+shares most of its prompt frees little. Offload helps a pool that runs short now and then.
+It doesn't replace the cap.
 
 **Single GPU.** Everything on rank 0:
 

mstar/communication/arena.py

@@ -657,8 +657,7 @@ def _stage_one(
             # file read for exactly those).
             data = _serialize_tensor(t)
             path = self._shm_path(self.my_entity_id, uuid)
-            with self._create_shm_file(path) as f:
-                f.write(data)
+            self._write_shm_file(path, data)
             self._shm_files[uuid] = path
             self._arena_ts[uuid] = time.monotonic()
             self.tensor_store.set_metadata(uuid, mem_registered=True)
@@ -734,7 +733,7 @@ def register_for_send(
         self, request_id: Rid, tensor_infos: list[TensorPointerInfo],
         skip_cuda_sync: bool = False,
     ):
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         queued = False
         placements = _Placements.empty()
@@ -758,7 +757,7 @@ def register_for_send_uuids(
         uuid, and placement goes back as two columns rather than a whole
         descriptor. The whole batch stages inside one stream context and syncs
         once at the end, rather than one host-blocking sync per request."""
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         queued = False
         placements = _Placements.empty()

mstar/communication/tensors.py

@@ -1,7 +1,9 @@
+import errno
 import hashlib
 import logging
 import os
 import platform
+import shutil
 import stat
 import time
 from abc import ABC, abstractmethod
@@ -133,30 +135,30 @@ def __init__(
         self.max_batch_size = max_batch_size
         self._executor = ThreadPoolExecutor(max_workers=max_workers)
         self._pending: list[Future] = []
-        if device != "cpu":
-            self._copy_stream = torch.cuda.Stream(device=device)
-        else:
-            self._copy_stream = torch.cuda.Stream()
+        # none on the host: host buffers have no GPU work to wait for, and a stream or event would
+        # create a CUDA context on a GPU that other processes may have filled
+        self._copy_stream = torch.cuda.Stream(device=device) if torch.device(device).type == "cuda" else None
 
     def submit(self, read_info: list[TransferReadInfo]) -> Future:
         """Non-blocking: enqueue a batch of READs.
 
-        Records a CUDA event on the current stream to ensure GPU data
+        On a device, records a CUDA event on the current stream to ensure GPU data
         is ready before the background thread reads it.
         """
         if not read_info:
             return
-        event = torch.cuda.current_stream().record_event()
+        event = torch.cuda.current_stream().record_event() if self._copy_stream is not None else None
         future = self._executor.submit(self._do_read, read_info, event)
         self._pending.append(future)
         # Prune completed futures to avoid unbounded growth
         self._pending = [f for f in self._pending if not f.done()]
         return future
 
-    def _do_read(self, read_info: list["TransferReadInfo"], event: torch.cuda.Event):
+    def _do_read(self, read_info: list["TransferReadInfo"], event: torch.cuda.Event | None):
         """Worker thread: wait for GPU data via CUDA event, then PUT."""
-        self._copy_stream.wait_event(event)
-        self._copy_stream.synchronize()
+        if event is not None:
+            self._copy_stream.wait_event(event)
+            self._copy_stream.synchronize()
 
         start_time = time.perf_counter()
 
@@ -327,6 +329,9 @@ def __init__(
         self.my_session_id = my_session_id
         self.enable_prof = enable_prof
         self.device = device
+        # by the manager's device, not torch.cuda.is_available() alone: a host manager holds only host tensors,
+        # and its first sync would create a CUDA context on a GPU that other processes may have filled
+        self._on_cuda = torch.cuda.is_available() and torch.device(device).type == "cuda"
         self.communicator = communicator
         self.transfer_engine = transfer_engine
         self.tensor_store = TensorStore()
@@ -464,7 +469,7 @@ def store_and_return_tensor_info(
         # step), so the unconditional default-stream sync here would
         # uselessly drain GPU(N+1). Pass ``skip_cuda_sync=True`` from
         # those call sites.
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         request_ids: list[Rid] = []
         uuids: list[int] = []
@@ -502,7 +507,7 @@ def store_and_return_tensor_info_batch(
         will ever read them, and a stored tensor nothing references is only
         freed when its request is torn down.
         """
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
 
         flat_uuids: list[int] = []
@@ -824,8 +829,13 @@ def get_ready_tensors(self, graph_walk: str | None=None) -> dict[Rid, list[Graph
         ready: dict[Rid, list[GraphEdge]] = {}
         still_pending = []
         uuid_to_time = {}
+        # the async reader finishes a request's reads in any order, so a read
+        # that is done waits for the ones started before it: a streamed token
+        # must not reach the client ahead of the one before it
+        waiting: set[Rid] = set()
         for ep in self.pending:
-            if ep.future is None or ep.future.done():
+            done = ep.future is None or ep.future.done()
+            if done and ep.request_id not in waiting:
                 if ep.future is not None:
                     ep.future.result()
                 for edge in ep.graph_edges:
@@ -838,6 +848,7 @@ def get_ready_tensors(self, graph_walk: str | None=None) -> dict[Rid, list[Graph
                         for info in edge.tensor_info:
                             uuid_to_time[info.uuid] = ep.rx_time
             else:
+                waiting.add(ep.request_id)
                 still_pending.append(ep)
         self.pending = still_pending
 
@@ -1127,7 +1138,7 @@ def __init__(
         )
 
     def register_for_send(self, request_id, tensor_infos, skip_cuda_sync=False):
-        if not skip_cuda_sync:
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         for info in tensor_infos:
             self._register_one(request_id, info.uuid)
@@ -1140,7 +1151,7 @@ def register_for_send_uuids(
         """Uuid-driven form. Registration here reads nothing off a descriptor
         but the uuid -- the tensor itself comes from the store -- so taking
         uuids avoids rebuilding one per tensor."""
-        if not skip_cuda_sync:
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         for request_id, uuids in per_request:
             for uuid in uuids:
@@ -1288,6 +1299,10 @@ def _deserialize_tensor(
     return t
 
 
+# Docker gives a container 64 MB of /dev/shm, and one video or audio tensor can outgrow that
+_SHM_DIR_WARN_BYTES = 1 << 30
+
+
 def _default_shm_dir() -> str:
     """Return the default shared-memory directory for the current platform."""
     if platform.system() == "Linux" and os.path.isdir("/dev/shm"):
@@ -1346,6 +1361,17 @@ def __init__(
 
         self.shm_dir = shm_dir or _default_shm_dir()
         os.makedirs(self.shm_dir, exist_ok=True)
+        # at boot, not at the first large send: that failure would be the only other sign of a small dir
+        try:
+            st = os.statvfs(self.shm_dir)
+            shm_total = st.f_frsize * st.f_blocks
+        except OSError:
+            shm_total = None
+        if shm_total is not None and shm_total < _SHM_DIR_WARN_BYTES:
+            logger.warning(
+                "SHM: %s holds only %d MiB and every tensor in flight is written there, so a large one will "
+                "fail its send with ENOSPC. Give it more room (docker --shm-size), or pass "
+                "--tensor-comm-protocol TCP, which is slower", self.shm_dir, shm_total >> 20)
 
         # uuid → file path for sender-side cleanup
         self._shm_files: dict[int, str] = {}
@@ -1404,11 +1430,40 @@ def _create_shm_file(self, path: str):
                     "deployment's namespace"
                 ) from None
 
+    def _write_shm_file(self, path: str, data: bytes) -> None:
+        """Write one tensor's bytes to ``path``.
+
+        A full shm dir fails the write part way, and no cleanup knows the partial file: left behind, it
+        would keep the space it took, and the next request would find the dir fuller still.
+        """
+        try:
+            with self._create_shm_file(path) as f:
+                f.write(data)
+        except OSError as exc:
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+            if exc.errno != errno.ENOSPC:
+                raise
+            free = shutil.disk_usage(self.shm_dir).free
+            if free < len(data):
+                reason = f"{self.shm_dir} has {free} bytes free, too few for a {len(data)}-byte tensor"
+            else:
+                reason = (
+                    f"{self.shm_dir} reported no space for a {len(data)}-byte tensor, with {free} bytes free after "
+                    "the failed write: it may be out of inodes, or other writers held the space meanwhile"
+                )
+            raise OSError(
+                errno.ENOSPC,
+                f"{reason}. Give it more room (docker --shm-size), or pass --tensor-comm-protocol TCP, which is slower",
+            ) from exc
+
     def register_for_send(
         self, request_id: Rid, tensor_infos: list[TensorPointerInfo],
         skip_cuda_sync: bool = False,
     ):
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         # Producer's completion event is already waited on upstream, so the
         # source tensors are device-visible on entry. Running the D2H on a
@@ -1430,7 +1485,7 @@ def register_for_send_uuids(
     ):
         """Uuid-driven form. This path reads nothing off a descriptor but the
         uuid, so taking uuids avoids rebuilding one per tensor."""
-        if not skip_cuda_sync and torch.cuda.is_available():
+        if not skip_cuda_sync and self._on_cuda:
             torch.cuda.default_stream().synchronize()
         ctx = (
             torch.cuda.stream(self._d2h_stream)
@@ -1451,8 +1506,7 @@ def _write_one(self, request_id: Rid, uuid: int) -> None:
         t0 = time.perf_counter()
         data = _serialize_tensor(tensor)
         path = self._shm_path(self.my_entity_id, uuid)
-        with self._create_shm_file(path) as f:
-            f.write(data)
+        self._write_shm_file(path, data)
         self._shm_files[uuid] = path
         self.tensor_store.set_metadata(uuid, mem_registered=True)
         if self.enable_prof:

mstar/engine/engine.py

@@ -595,8 +595,10 @@ def _open_prefix_caches(self, specs_by_key, model) -> None:
                 logger.warning(
                     "KV %s: prefix cache on with cpu_offload_pages 0: a decode "
                     "step that finds nothing to evict holds its requests until "
-                    "they time out. max_concurrent_requests caps how many run "
-                    "at once, and cpu_offload_pages gives the worker a victim",
+                    "they time out. Set max_concurrent_requests to the pool's "
+                    "pages over the pages one request needs. cpu_offload_pages "
+                    "gives the worker a victim, but one that shares most of its "
+                    "prompt frees little",
                     key,
                 )
 

mstar/model/bagel/bagel_model.py

@@ -736,14 +736,14 @@ def _kv_config(self) -> PagedKVConfig:
         )
 
     def prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]:
-        """The text walk's prompt is its token ids, and so is every step it takes.
+        """The text walk's prompt is its token ids.
 
         The image walks write into this stream too, and an edit's picture lands
         at the positions its keyed text would have; only the text walk is named.
+        Generated pages are not keyed: the chat API sends a reply back inside
+        the next user turn, so their keys never match.
         """
-        return {"kv": {"main": PrefixStream(
-            "text_inputs", "ids", "prefill_text", "decode",
-        )}}
+        return {"kv": {"main": PrefixStream("text_inputs", "ids", "prefill_text")}}
 
     def checkpoint_path(self) -> str:
         return snapshot_download(

mstar/model/orpheus/orpheus_model.py

@@ -359,10 +359,9 @@ def get_initial_forward_pass_args(
     # -------------------------------------------------------------------
 
     def prefix_key_streams(self) -> dict[str, dict[str, PrefixStream]]:
-        """The LLM's prompt is its token ids, and so is every step after it."""
-        return {
-            KV_CACHE: {"main": PrefixStream("text_inputs", "ids", "prefill", "decode")},
-        }
+        """The LLM's prompt is its token ids. Its generated audio codes never
+        come back in a prompt, so they are not keyed."""
+        return {KV_CACHE: {"main": PrefixStream("text_inputs", "ids", "prefill")}}
 
     def checkpoint_path(self) -> str:
         return _resolve_local_hf_snapshot(