Host overhead 2/N: API server output path, one data-worker item per frame, raw token bytes, batched Rust-frontend relay

docs/adding_models.rst

@@ -166,10 +166,6 @@ You must implement these abstract methods:
    (``request_done=True``). A simple prefill-decode model sets ``is_prefill`` to false
    once, then repeats the decode walk until EOS.
 
-``postprocess(self, output, modality) -> bytes``
-   Encode a finished output tensor to bytes for the client: ``utf-8`` for text, PNG for
-   images, raw PCM for audio, and so on.
-
 ``get_submodule(self, node_name, device="cpu", tp_group=None, autocast_dtype=None, sp_group=None, moe_parallel=None) -> NodeSubmodule | None``
    Build and return the ``NodeSubmodule`` for ``node_name``, and cache the result. Load
    weights here, on ``device``. Return ``None`` for dummy mode.
@@ -226,9 +222,32 @@ You must implement these abstract methods:
    ``kwargs["request"]``, since its stage-1 walk sends a half-size ``height``.
 
 The following methods have defaults and are optional to override:
+``text_output``, ``postprocess`` and ``postprocess_inline`` (described next),
 ``get_request_resource_configs`` (described below), ``get_sampling_config``,
 ``get_max_output_tokens``, ``get_autocast_dtype``, ``load_image``, ``load_audio``,
 ``load_video``, and the partition methods described at the end of this page.
+
+``text_output(self) -> TextOutput | None``
+   Declare how text output (token ids) becomes bytes: a detokenizer, and the ids that
+   end a reply, which are dropped unless the request sets ``ignore_eos``.
+   ``postprocess`` and ``postprocess_inline`` derive their text paths from it. A text
+   chunk may end inside a character, since the API server joins a request's text
+   chunks before decoding them, so a byte-level BPE model should declare
+   ``mstar.model.utils.TextChunkDetokenizer`` (or ``ByteLevelDetokenizer``), which
+   emits each token's raw bytes, rather than ``tokenizer.decode([id])``, which turns
+   a split character into U+FFFD.
+
+``postprocess(self, output, modality, request_kwargs=None) -> bytes``
+   Encode a finished output tensor to bytes for the client: ``utf-8`` for text, PNG for
+   images, raw PCM for audio, and so on. Override it for any modality other than
+   text, and hand text to ``super().postprocess``.
+
+``postprocess_inline(self, values, modality, request_kwargs) -> list[bytes]``
+   ``postprocess`` for one decode step's inline values (one sampled token id per
+   request), called once per step. With ``text_output`` declared, each id maps to
+   its memoized bytes; otherwise each value goes through ``postprocess`` as a
+   one-element tensor.
+
 ``Model.nodes`` is a read-only property. It returns the sorted set of node names that
 appear in any graph walk.
 

docs/installation.rst

@@ -554,6 +554,11 @@ it. Then opt in:
    mstar serve qwen3_omni --rust-frontend
    # or, with an explicit config: mstar-serve --config <model.yaml> --rust-frontend
 
+Use it when benchmarking or tuning streaming throughput: its output path costs
+about 8 µs of API-server CPU per streamed token against 18 µs for the Python
+frontend (``benchmark/api_output_path``), so at high token rates the Python
+frontend, not the model, can be the limit.
+
 The binary is resolved via ``--rust-frontend-bin`` / ``MSTAR_SERVER_BIN`` /
 the environment's scripts directory / ``$PATH`` / the in-repo build. Its knobs
 are listed in :doc:`environment_variables`.

mstar/api_server/data_worker.py

@@ -159,12 +159,13 @@ def _video_frame_metadata(
         "frame_count": frame_count,
     }
 @dataclass
-class InlineToken:
-    """One request's value out of a RESULT_TOKENS frame, queued behind the
-    tensor results so sequences keep arrival order."""
-    request_id: str
-    value: int
-    loop_indices: NestedLoopIndices | None
+class InlineTokens:
+    """One RESULT_TOKENS frame's values for the requests still known here,
+    queued behind the tensor results so each request's chunks keep arrival
+    order. One item per frame, not per request. No loop indices: a chunk's
+    order is its arrival sequence, and nothing reads them."""
+    request_ids: list[str]
+    values: list[int]
     modality: str
 
 
@@ -363,31 +364,29 @@ def new_result_tensors(self, input: ResultTensors):
         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,
+    def new_result_tokens(
+        self, request_ids: list[str], values: list[int], signal: str, modality: str,
     ):
-        """A value that rode inline: same accounting as a result tensor, no
-        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.received_chunk_counts:
-            logger.debug("Late token for cleaned-up request %s dropped", request_id)
-            return
-        counts = self.received_chunk_counts[request_id]
-        counts[signal] = counts.get(signal, 0) + 1
-        self.per_request_reading_tensors[request_id] += 1
-        self.result_tensor_input_queue.put(
-            InlineToken(request_id, value, loop_indices, modality)
-        )
-        if wake:
+        """A frame's values that rode inline: the same accounting per request
+        as a result tensor, no transport read and nothing to ack. A request
+        already cleaned up is dropped."""
+        counts_by_rid = self.received_chunk_counts
+        reading = self.per_request_reading_tensors
+        kept_ids: list[str] = []
+        kept_values: list[int] = []
+        for rid, value in zip(request_ids, values, strict=True):
+            counts = counts_by_rid.get(rid)
+            if counts is None:
+                logger.debug("Late token for cleaned-up request %s dropped", rid)
+                continue
+            counts[signal] = counts.get(signal, 0) + 1
+            reading[rid] += 1
+            kept_ids.append(rid)
+            kept_values.append(value)
+        if kept_ids:
+            self.result_tensor_input_queue.put(InlineTokens(kept_ids, kept_values, modality))
             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.
 
@@ -1196,35 +1195,48 @@ def _read_result_tensor(
             state.order[tensor_info.uuid] = (state.next_sequence, result.loop_indices)
             state.next_sequence += 1
 
-    def _emit_inline_token(self, item: InlineToken) -> None:
-        """Detokenize a value that arrived inline and queue its chunk."""
-        request_id = item.request_id
-        state = self.request_output_state.setdefault(
-            request_id, RequestOutputState()
-        )
-        sequence = state.next_sequence
-        state.next_sequence += 1
+    def _emit_inline_tokens(self, item: InlineTokens) -> None:
+        """Postprocess a frame's inline values in one model call and queue each
+        request's chunk, in arrival order with its tensor results."""
+        rids, values = item.request_ids, item.values
+        if self._draining_rids:
+            keep = [i for i, rid in enumerate(rids) if rid not in self._draining_rids]
+            rids = [rids[i] for i in keep]
+            values = [values[i] for i in keep]
+        kwargs = self.request_model_kwargs
         try:
-            tensor = torch.tensor([item.value], dtype=torch.int64)
-            postprocessed = self.model.postprocess(
-                tensor, item.modality,
-                request_kwargs=self.request_model_kwargs.get(request_id),
-            )
-            self._queue_completed_output(
-                request_id,
-                sequence,
-                ResultChunk(
-                    request_id=request_id,
-                    modality=item.modality,
-                    data=postprocessed,
-                    metadata={},
-                ),
-            )
-        except Exception as exc:  # noqa: BLE001 — must reach the client
-            self._fail_request(
-                request_id, exc, f"{item.modality} output postprocessing",
-                sequence=sequence,
+            datas = self.model.postprocess_inline(
+                values, item.modality, [kwargs.get(rid) for rid in rids],
             )
+        except Exception:  # noqa: BLE001 — retried row by row below
+            # so a failure fails only the request it belongs to
+            datas = None
+        states = self.request_output_state
+        for i, request_id in enumerate(rids):
+            state = states.get(request_id)
+            if state is None:
+                state = states[request_id] = RequestOutputState()
+            sequence = state.next_sequence
+            state.next_sequence += 1
+            try:
+                if datas is None:
+                    data = self.model.postprocess_inline(
+                        [values[i]], item.modality, [kwargs.get(request_id)],
+                    )[0]
+                else:
+                    data = datas[i]
+                self._queue_completed_output(
+                    request_id,
+                    sequence,
+                    ResultChunk(
+                        request_id=request_id, modality=item.modality, data=data, metadata={},
+                    ),
+                )
+            except Exception as exc:  # noqa: BLE001 — must reach the client
+                self._fail_request(
+                    request_id, exc, f"{item.modality} output postprocessing",
+                    sequence=sequence,
+                )
 
     def _queue_completed_output(
         self,
@@ -1656,11 +1668,10 @@ def run(self):
                 while not self.result_tensor_queue.empty():
                     did_work = True
                     result = self.result_tensor_queue.get()
-                    if isinstance(result, InlineToken):
+                    if isinstance(result, InlineTokens):
                         # Nothing to read or ack: postprocess in place, in
                         # arrival order with the tensor results.
-                        if result.request_id not in self._draining_rids:
-                            self._emit_inline_token(result)
+                        self._emit_inline_tokens(result)
                         continue
                     # Draining for teardown: don't start new reads — ack the
                     # tensors back so the producing worker can free its buffers.