worker: device loop-back, rows-only steps, leaner API delivery (round 2)

docs/environment_variables.rst

@@ -236,6 +236,13 @@ Read by the ``mstar-server`` binary and its bridge
    * - ``MSTAR_MAX_BODY_MB``
      - ``128``
      - Request body limit (multipart uploads included).
+   * - ``MSTAR_TOKENIZE_PARALLEL``
+     - ``1``
+     - When a long prompt is encoded as parallel chunks, set
+       ``TOKENIZERS_PARALLELISM=true`` for the process first: the Rust
+       tokenizer only uses its thread pool while that variable reads true,
+       and serving environments often export it false. ``0`` leaves the
+       variable as it is (the chunked encode then runs serially).
    * - ``MSTAR_TOKENIZER``
      - unset
      - Path to a HuggingFace ``tokenizer.json`` enabling frontend
@@ -315,6 +322,32 @@ Worker scheduling
      - ``1024``
      - Cap on back-to-back speculative steps before the leader yields to
        other ready work.
+   * - ``MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS``
+     - ``32,64,128,256,512,1024,2048,4096,8192``
+     - Total tokens of a packed Qwen3.5 text-prefill step that get a captured
+       graph (one per bucket and batch size). A prompt past the top bucket
+       prefills eagerly, which is launch-bound: an eager 8k prefill on 0.8B
+       spends 78 ms of wall time on 36 ms of kernels. The default now reaches
+       8192; narrow it to save graph memory and capture time, or extend it.
+   * - ``MSTAR_DEVICE_LOOPBACK``
+     - ``1``
+     - Whether a decode node may keep its loop-back token on the device:
+       the sampler writes each request's last token to a slot-addressed
+       master and the next step reads its input ids from it, so the worker
+       routes the signal with no tensor (no per-request store, hold or
+       cleanup per step). Only nodes that opt in (Qwen3.5 decode) and only
+       when the sampler has graph buffers; ``0`` keeps every node on the
+       per-request tensor path.
+   * - ``MSTAR_SAMPLER_INGRAPH_SCATTER``
+     - ``0``
+     - The captured sampler gathers its RNG offsets from the slot masters
+       and scatters the advanced offsets and the sampled tokens back inside
+       the graph, so a replayed decode step issues none of those launches on
+       the gpu thread afterwards (each one is a GIL release the main thread's
+       postprocess wins). Padding rows then address a trash master row and
+       the masters are allocated once for ``MSTAR_SAMPLER_SLOTS`` concurrent
+       requests (default 4096, a hard cap with this on). Only with the device
+       loop-back on. An experiment knob; measure before turning it on.
    * - ``MSTAR_RUST_SEND_TIMING``
      - unset
      - ``1`` makes the Rust runtime print, every 500 calls, how long the

mstar/api_server/data_worker.py

@@ -25,6 +25,7 @@
     ResultTensors,
 )
 from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator
+from mstar.communication.event import EventWakeup
 from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager
 from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig
 from mstar.engine.resources.kv.keys import chain
@@ -138,6 +139,8 @@ def __init__(
         self.output_queue = queue.Queue()
         self.profile_queue = queue.Queue()
         self.stop_event = threading.Event()
+        # Wakes the thread's idle wait: every producer below signals it.
+        self.wakeup = EventWakeup()
 
         self.per_request_reading_tensors = {}
         self.output_loop_idxs: dict[str, NameToLoopIndices] = {}
@@ -174,6 +177,7 @@ def __init__(
                 reads_done_queue=self.reads_done_queue,
                 discard_tensor_queue=self.discard_tensor_queue,
                 stop_event=self.stop_event,
+                wakeup=self.wakeup,
                 communicator=self.communicator,
                 tensor_manager=self.tensor_manager,
                 model=model,
@@ -184,10 +188,18 @@ def __init__(
         )
         self.thread.start()
 
+    def _signal(self):
+        """End the thread's idle wait: new work is queued. getattr: tests
+        build this facade without its constructor."""
+        wakeup = getattr(self, "wakeup", None)
+        if wakeup is not None:
+            wakeup.signal()
+
     def new_request(self, input: PreprocessInput):
         self.output_loop_idxs[input.request_id] = {}
         self.per_request_reading_tensors[input.request_id] = 0
         self.request_input_queue.put(input)
+        self._signal()
 
     def abort_request(self, request_id: str):
         # Forward the abort to the conductor and begin draining our own reads
@@ -196,6 +208,7 @@ def abort_request(self, request_id: str):
         # REMOVE_REQUEST, so a worker still reading them can't be unlinked
         # out from under it. Drop only the main-thread bookkeeping here.
         self.abort_request_queue.put(request_id)
+        self._signal()
         self.output_loop_idxs.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
@@ -207,6 +220,7 @@ def finished_reading(self, request_id: str, drained: bool = True):
         when delivery was abandoned (TTL, client gone) and the ACK must wait
         for any in-flight read."""
         self.reads_done_queue.put((request_id, drained))
+        self._signal()
         self.output_loop_idxs.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
@@ -229,13 +243,17 @@ def new_result_tensors(self, input: ResultTensors):
             input.request_id,  self.per_request_reading_tensors[input.request_id]
         )
         self.result_tensor_input_queue.put(input)
+        self._signal()
 
     def new_result_token(
         self, request_id: str, value: int,
         loop_indices: NestedLoopIndices | None, signal: str, modality: str,
+        wake: bool = True,
     ):
         """A value that rode inline: same accounting as a result tensor, no
-        transport read and nothing to ack."""
+        transport read and nothing to ack. ``wake=False`` queues without
+        ending the thread's idle wait: a frame's caller hands over every
+        request's value first and calls ``wake`` once."""
         if request_id not in self.output_loop_idxs:
             logger.debug("Late token for cleaned-up request %s dropped", request_id)
             return
@@ -247,6 +265,12 @@ def new_result_token(
         self.result_tensor_input_queue.put(
             InlineToken(request_id, value, loop_indices, modality)
         )
+        if wake:
+            self._signal()
+
+    def wake(self):
+        """End the thread's idle wait after a batch of ``wake=False`` puts."""
+        self._signal()
 
     def discard_result_tensors(self, input: ResultTensors):
         """Ack and drop result tensors for an already-removed request.
@@ -255,6 +279,7 @@ def discard_result_tensors(self, input: ResultTensors):
         producing worker gets its TENSOR_RECEIVED ack and frees the buffers.
         """
         self.discard_tensor_queue.put(input)
+        self._signal()
 
     def has_pending_tensors(self, request_id: str):
         return self.per_request_reading_tensors.get(request_id, 0) > 0
@@ -324,11 +349,13 @@ def get_rx_info(self, request_id: str) -> list[RxInfo]:
 
     def cleanup_request(self, request_id: str):
         self.cleanup_request_queue.put(request_id)
+        self._signal()
         self.output_loop_idxs.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
     def shutdown(self):
         self.stop_event.set()
+        self._signal()
         if self.thread.is_alive():
             self.thread.join()
 
@@ -384,6 +411,7 @@ def __init__(
         enable_prof: bool=False,
         model_config: dict | None = None,
         enable_nvtx: bool=False,
+        wakeup: EventWakeup | None = None,
     ):
         # keying a stream needs the deployment's page size as well as the
         # model's declaration, so a worker built without a config keys no stream
@@ -428,6 +456,12 @@ def __init__(
         # Owned by PreprocessWorker (main thread); used only from this thread.
         self.communicator = communicator
         self.tensor_manager = tensor_manager
+        # Set by the producers (queue puts, shutdown) and by finished tensor
+        # reads; run() blocks on the communicator's poll with it instead of
+        # sleeping and sweeping the queues. None keeps the sleep (tests built
+        # with a bare communicator).
+        self.wakeup = wakeup
+        self.idle_wait_ms = 20
 
     def _cleanup_request_state(self, request_id: str, *, force: bool = False) -> None:
         """Release transport and postprocessing state owned by this thread.
@@ -677,10 +711,14 @@ def _read_result_tensor(
         self, result: ResultTensors
     ):
         result.graph_edge.name = f"{result.modality}_output"
-        self.tensor_manager.start_read_tensors(
+        futures = self.tensor_manager.start_read_tensors(
             request_id=result.request_id,
             graph_edges=[result.graph_edge],
         )
+        wakeup = getattr(self, "wakeup", None)
+        if wakeup is not None and futures:
+            # a finished read ends the idle wait, so the chunk goes out at once
+            wakeup.register_futures(futures)
         if result.request_id not in self.tensor_uuid_to_metadata_per_request:
             self.tensor_uuid_to_metadata_per_request[result.request_id] = {}
         state = self.request_output_state.setdefault(
@@ -930,11 +968,28 @@ def _hard_cleanup(self, request_id: str) -> None:
         self.tensor_manager.force_cleanup_request(request_id)
         self._drop_request_state(request_id)
 
+    def _idle_wait(self):
+        """How a pass that found nothing waits: on the communicator's poll,
+        ended by a message, a producer's signal or a finished read, with
+        ``idle_wait_ms`` as the backstop; or a 1 ms sleep when the thread has
+        no wakeup or its communicator cannot poll one."""
+        # getattr: teardown tests build a thread without its constructor
+        wakeup = getattr(self, "wakeup", None)
+        comm = getattr(self, "communicator", None)
+        register = getattr(comm, "register_event_for_poll", None)
+        wait = getattr(comm, "wait_for_work", None)
+        if wakeup is None or not callable(register) or not callable(wait):
+            return lambda: time.sleep(0.001)
+        register(wakeup)
+        idle_ms = getattr(self, "idle_wait_ms", 20)
+        return lambda: wait(idle_ms)
+
     def run(self):
         # A thread told to stop before it starts (teardown tests build one with
         # only its queues) has nothing to warm up.
         if not self.stop_event.is_set():
             warm_up_model(getattr(self, "model", None))
+        idle_wait = self._idle_wait()
         while not self.stop_event.is_set():
             did_work = False
             try:
@@ -1018,7 +1073,7 @@ def run(self):
                 logger.exception("PreprocessWorkerThread error")
 
             if not did_work:
-                time.sleep(0.001)
+                idle_wait()
 
         # Stopping, and nothing will send the RemoveRequest for what is still
         # tracked (the conductor is gone or going), so drop it here rather than

mstar/api_server/entrypoint.py

@@ -34,6 +34,7 @@
 from mstar.profile.format import OutputInfo, RequestProfile, RequestTiming
 from mstar.utils import profiler
 from mstar.utils.exitcode import describe_exitcode
+from mstar.utils.gc_freeze import freeze_after_setup
 from mstar.utils.logging_config import quiet_noisy_loggers
 from mstar.utils.orphan import watch_parent
 
@@ -366,6 +367,9 @@ def finalize_setup(self) -> None:
                     and message.message_type == "setup_done"
                 ):
                     logger.info("All workers ready")
+                    # the set-up heap is complete: keep gen-2 collections off
+                    # it (measured 200 ms per pass on this process under load)
+                    freeze_after_setup("API server")
                     self._msg_thread.start()
                     return
                 logger.warning(
@@ -578,18 +582,24 @@ def _process_messages(self) -> None:
                         # One frame for the step: every request's value.
                         # A request that is gone needs nothing acked back.
                         body = message.body
+                        queued = False
                         with self.request_lock:
                             for i, rid in enumerate(body.request_ids):
                                 if rid in self.pending_requests:
+                                    # one wake-up for the frame, after every
+                                    # request's value is queued
                                     self.preprocess_worker.new_result_token(
                                         rid, body.values[i],
                                         body.loop_indices[i],
-                                        body.signal, body.modality,
+                                        body.signal, body.modality, wake=False,
                                     )
+                                    queued = True
                                 elif rid not in self.recently_completed:
                                     logger.warning(
                                         "Token for unknown request %s dropped", rid,
                                     )
+                        if queued:
+                            self.preprocess_worker.wake()
                         continue
 
                     rid = message.body.request_id
@@ -733,14 +743,31 @@ def _process_messages(self) -> None:
     # ----------------------------------------------------------
 
     async def iter_result_chunks(self, request_id: str):
-        """Yield raw :class:`ResultChunk` objects as they arrive.
+        """Yield raw :class:`ResultChunk` objects as they arrive, one at a time.
 
         Shared source for both output surfaces: ``/generate`` formats each
         chunk as NDJSON (via :meth:`async_stream_results`), while the
-        OpenAI-compatible endpoints translate the same chunks into SSE. The
-        per-request timeout, incremental drain, and final flush behave exactly
-        as before — only the yielded type changed (``ResultChunk`` instead of a
-        pre-serialized line).
+        OpenAI-compatible endpoints translate the same chunks into SSE. This
+        is a flattening wrapper over :meth:`iter_result_chunk_batches`, which
+        hands out the chunks the way they arrive (several per wake under load)
+        and owns the timeout, the drain and the abort on an early exit.
+        """
+        batches = self.iter_result_chunk_batches(request_id)
+        try:
+            async for batch in batches:
+                for chunk in batch:
+                    yield chunk
+        finally:
+            await batches.aclose()
+
+    async def iter_result_chunk_batches(self, request_id: str):
+        """Yield the :class:`ResultChunk` objects as they arrive, as one list
+        per drain: everything the message thread appended since the last
+        wake. Under load several tokens of a request are often waiting when
+        the handler runs, and sending them as one body frame costs one write
+        instead of one per token; when the loop keeps up the lists are of
+        length one and nothing is delayed. The final flush and an in-band
+        error come as the last list.
         """
         start = time.time()
         finished = False
@@ -773,10 +800,11 @@ async def iter_result_chunks(self, request_id: str):
                     else:
                         done = True
 
-                for chunk in new_chunks:
+                if new_chunks:
                     if self.enable_nvtx:
-                        profiler.mark("apiserver.chunk_available")
-                    yield chunk
+                        for _ in new_chunks:
+                            profiler.mark("apiserver.chunk_available")
+                    yield new_chunks
 
                 if done:
                     logger.info("Async stream results received finish for %s", request_id)
@@ -792,19 +820,20 @@ async def iter_result_chunks(self, request_id: str):
                     # lock; the popped request is no longer shared with other threads.
                     if finished_req is not None:
                         self._finalize_profile(finished_req)
-                    for chunk in remaining:
-                        yield chunk
+                    tail = list(remaining)
                     # 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:
-                        yield ResultChunk(
+                        tail.append(ResultChunk(
                             request_id=request_id,
                             modality="error",
                             data=str(finished_req.error).encode("utf-8"),
                             metadata={"status": finished_req.error_status},
-                        )
+                        ))
+                    if tail:
+                        yield tail
                     finished = True
                     break
 

mstar/api_server/openai/_util.py

@@ -18,7 +18,18 @@ def rid(prefix: str) -> str:
 
 
 def sse(obj: dict) -> str:
-    return f"data: {json.dumps(obj)}\n\n"
+    return f"data: {_dumps(obj)}\n\n"
+
+
+try:
+    import orjson
+
+    def _dumps(obj: dict) -> str:
+        # ~5x json.dumps; one event per token at tens of thousands a second
+        return orjson.dumps(obj).decode()
+except ImportError:  # pragma: no cover
+    def _dumps(obj: dict) -> str:
+        return json.dumps(obj)
 
 
 def error_type(status: int) -> str: