[Chunked Prefill #6] Arrival-order batch builder

docs/adding_models.rst

@@ -1142,7 +1142,10 @@ chunk is one iteration of a Loop around the node. To opt a walk in:
   - ``max_batch_tokens``: the most tokens a step runs. A combined walk that carries
     chunked rows sets one too, since its steps run under it.
   - ``accumulated_outputs``: output edges a consumer needs every chunk of. They arrive as
-    one item per chunk. Any other edge carries the final chunk's output alone.
+    one item per chunk, all when the prompt finishes.
+  - ``per_chunk_outputs``: streaming edges sent on every chunk as it runs, for a consumer
+    that can start on a prefix. Any edge in neither set carries the final chunk's output
+    alone.
   - ``reuses_cached_prefix``: a callable saying whether a row may start past its cached
     prefix; ``None`` (the default) is always. Set it where a row may have to write its
     prefix again.
@@ -1153,8 +1156,11 @@ chunk is one iteration of a Loop around the node. To opt a walk in:
 
   A streaming consumer of an accumulated edge whose ``ChunkPolicy`` expects one item
   (``FixedChunkPolicy(1)``, say) sets ``Connection.join_dim``, so its ``StreamBuffer``
-  joins the pieces into that item. Leave it unset for a consumer that can take the
-  pieces as they come.
+  joins the pieces into that item. A consumer that can take the pieces as they come says
+  so with ``ChunkPolicy.allow_partial_input`` (``FixedChunkPolicy(allow_partial=True)``),
+  and the producer then sends the edge per chunk, as Qwen3-Omni's Thinker does into its
+  Talker. A chunked stream its consumer would misread, each piece taken for the whole
+  input, is refused at startup.
 - ``get_input_sequence_len(graph_walk, fwd_info, inputs)`` gives a row's length from its
   input shapes, before ``prepare_inputs`` runs. The scheduler cuts the row by it, and a
   row whose prepared input turns out another length fails.

docs/serving.rst

@@ -309,7 +309,12 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
        a prompt mid-prefill takes the budget before a newer one, so it finishes first,
        and prompts short enough for an equal share still go in whole. ``fair_share``
        splits the budget max-min fair over every waiting prompt, so each advances every
-       step and all finish late together.
+       step and all finish late together. ``arrival`` serves prompts strictly in arrival
+       order: the first that does not fit is cut and every prompt behind it waits, a
+       short one too. It also keeps one prompt per node partway through: while that
+       prompt is outside a step (held after an OOM, or its chunk still running), no
+       other prompt starts, which keeps prompts on a tight KV pool from each holding
+       pages the others need.
    * - ``min_chunk_tokens``
      - *(optional, default 128)* The fewest tokens a prompt cut by chunked prefill runs
        in a step, or the whole budget if that is smaller. A positive int, for every

mstar/engine/chunking.py

@@ -104,7 +104,7 @@ def set_budgets(
         of such a walk would do nothing, so it is refused; the combined walk
         itself can be named instead.
 
-        ``no_chunking`` names nodes this worker chunks nowhere (TP/SP, which
+        ``no_chunking`` names nodes this worker chunks nowhere (SP, which
         chunked prefill does not support yet). A budget set on one is still
         validated but has no effect, since the deployment, not the config, is
         what turned chunking off there.

mstar/model/base.py

@@ -47,7 +47,9 @@ def _wrap_node_in_loop(
     The node's real edges move onto the Loop, which is what makes a chunk's
     output the loop's business rather than the worker's: an edge the policy
     gathers gets every chunk's tensors at the end, and any other gets the
-    final chunk's alone, earlier ones dropped as the loop advances.
+    final chunk's alone, earlier ones dropped as the loop advances. A
+    ``per_chunk_outputs`` edge is the exception: it stays on the body, so
+    every chunk routes it as it completes.
 
     The body keeps a bare same-named edge to ``EMPTY_DESTINATION``. The Loop
     caches by edge NAME, so the body has to emit the name for the Loop to have
@@ -65,6 +67,9 @@ def _wrap_node_in_loop(
             return section
         ticks, final, gathered = [], [], []
         for edge in section.outputs:
+            if edge.name in policy.per_chunk_outputs:
+                ticks.append(edge)
+                continue
             ticks.append(GraphEdge(next_node=EMPTY_DESTINATION, name=edge.name))
             dest = gathered if edge.name in policy.accumulated_outputs else final
             dest.append(edge)

mstar/model/qwen3_omni/qwen3_omni_model.py

@@ -21,7 +21,9 @@
     KV cache in ``talker_prefill``; the first ``thinker_decode`` state runs
     ``talker_last_prefill``, which samples the first codec token, and the
     rest feed ``talker_decode``.  The Talker never needs to know how many
-    prefill walks the Thinker runs.
+    prefill walks the Thinker runs.  A chunked Thinker prefill streams each
+    chunk's states as it runs (``per_chunk_outputs``), and the Talker extends
+    its KV with each, so it keeps pace with the Thinker's prompt.
 
 Text-only mode:
     When output_modalities does not include "audio", only the Thinker
@@ -124,9 +126,11 @@ def apply_code2wav_serving_overrides(code2wav, model_kwargs: dict) -> None:
 
 def _thinker_to_talker_policy() -> FixedChunkPolicy:
     # Past the Thinker's end only the decode loop keeps running; an empty
-    # chunk in a prefill walk would run a prefill pass with nothing in it
+    # chunk in a prefill walk would run a prefill pass with nothing in it.
+    # Partial: the Talker extends its KV with each Thinker prefill chunk.
     return FixedChunkPolicy(
         chunk_size=1, continue_after_done=frozenset({"talker_decode"}),
+        allow_partial=True,
     )
 
 
@@ -378,9 +382,8 @@ def get_chunking_policies(self) -> dict[str, tuple[ChunkingPolicy, ...]]:
             ChunkingPolicy(
                 walks=("prefill_text", "prefill_audio", "prefill_vision"),
                 chunkable=True, max_batch_tokens=budget,
-                # the Talker conditions on the whole prompt, so it needs every
-                # chunk's states, not just the last one's
-                accumulated_outputs=frozenset({"thinker_states", "thinker_mask"}),
+                # the Talker extends its KV with each chunk's states as it lands
+                per_chunk_outputs=frozenset({"thinker_states", "thinker_mask"}),
             ),
             # carries the budget for the steps that mix a chunk with decode rows
             ChunkingPolicy(walks=(THINKER_MIXED,), max_batch_tokens=budget),
@@ -650,16 +653,13 @@ def get_partition_topology(self) -> PartitionTopology:
                     edge_name="thinker_states",
                     chunk_policy_factory=_thinker_to_talker_policy,
                     consumer_walk=_talker_walk,
-                    # the Thinker's prefill chunks are pieces of one prompt
-                    join_dim=0,
                 ),
                 Connection(
                     from_partition="Thinker",
                     to_partition="Talker",
                     edge_name="thinker_mask",
                     chunk_policy_factory=_thinker_to_talker_policy,
                     consumer_walk=_talker_walk,
-                    join_dim=0,
                 ),
                 Connection(
                     from_partition="Talker",

mstar/model/qwen3_omni/submodules.py

@@ -1191,8 +1191,15 @@ def postprocess(
                 selected = torch.where(
                     multimodal_mask.unsqueeze(-1), layer_n, layer_0,
                 )                              # (seq_len, hidden)
-                outputs["thinker_states"] = [selected[inclusion_mask]]
-                outputs["thinker_mask"] = [multimodal_mask[inclusion_mask]]
+                kept = selected[inclusion_mask]
+                if kept.shape[0]:
+                    outputs["thinker_states"] = [kept]
+                    outputs["thinker_mask"] = [multimodal_mask[inclusion_mask]]
+                else:
+                    # A chunk of only system prompt or earlier turns gives the
+                    # Talker nothing to write; unsent, it runs no empty prefill.
+                    outputs.pop("thinker_states")
+                    outputs.pop("thinker_mask")
 
         if "new_token" not in outputs:
             return
@@ -1415,9 +1422,7 @@ def prepare_inputs(
         device = self.get_device()
 
         if graph_walk == "talker_prefill":
-            # One item, whether or not the Thinker's prefill was chunked: its
-            # per-chunk pieces are joined on arrival by the StreamBuffer
-            # (Connection.join_dim), so the Talker sees what it always did.
+            # The Thinker's walk, or one chunk of it when its prefill is chunked
             selected_hidden = inputs["thinker_states"][0].to(device)
             multimodal_mask = inputs["thinker_mask"][0]
             input_embeds = self._get_talker_embeds(

mstar/model/submodule_base.py

@@ -369,7 +369,12 @@ class ChunkingPolicy:
     ``accumulated_outputs``: output edges a consumer needs every chunk of,
     gathered on the Loop and handed over as one item per chunk when the
     prompt finishes; the consumer joins the pieces itself. An edge not named
-    here carries the final chunk's tensors alone.
+    here or in ``per_chunk_outputs`` carries the final chunk's tensors alone.
+
+    ``per_chunk_outputs``: streaming output edges sent on every chunk as it
+    runs, each chunk's piece an item of its own, for a consumer that can
+    start on a prefix (``ChunkPolicy.allow_partial_input``, checked at load).
+    Disjoint from ``accumulated_outputs``.
 
     ``max_chunks``: the chunk Loop's ``max_iters``, a runaway guard rather
     than the real count, which is not known until the budget cuts each step.
@@ -385,13 +390,21 @@ class ChunkingPolicy:
     # rows may be split: `split_inputs` cuts them, `get_input_sequence_len` measures them
     chunkable: bool = False
     accumulated_outputs: frozenset[str] = frozenset()
+    per_chunk_outputs: frozenset[str] = frozenset()
     max_chunks: int = 2**31 - 1
     # most tokens a step of this walk runs; a combined walk that carries chunks sets one too
     max_batch_tokens: int | None = None
     # whether a row may start past its cached prefix, asked before prepare; None is always
     reuses_cached_prefix: Callable[[CurrentForwardPassInfo], bool] | None = None
     reserve_whole_prefill: bool = True
 
+    def __post_init__(self):
+        both = self.accumulated_outputs & self.per_chunk_outputs
+        if both:
+            raise ValueError(
+                f"outputs {sorted(both)} are both gathered and sent per chunk; name each in one"
+            )
+
 
 NO_CHUNKING = ChunkingPolicy()
 

mstar/streaming/chunk_policy.py

@@ -46,6 +46,13 @@ def window_size(self) -> int:
         """
         ...
 
+    def allow_partial_input(self) -> bool:
+        """Whether the consumer takes a chunked prefill a chunk at a time, each
+        chunk's piece an item of its own (``ChunkingPolicy.per_chunk_outputs``).
+        Default ``False``: a stream that does so into this consumer is refused
+        at load, since the consumer would read a piece as the whole input."""
+        return False
+
     def continue_after_producer_done(self) -> bool:
         """Whether the buffer should keep producing (empty) chunks after the
         producer signals done and all buffered items have been consumed.
@@ -150,14 +157,18 @@ class FixedChunkPolicy(ChunkPolicy):
             the producer finishes and all buffered items are consumed. A set
             of consumer walks does so only while the consumer is in one of
             them; in its other walks the stream just waits.
+        allow_partial: if True, take a chunked prefill's pieces as they
+            land (``ChunkPolicy.allow_partial_input``).
     """
 
     def __init__(
         self, chunk_size: int, continue_after_done: bool | frozenset[str] = False,
+        allow_partial: bool = False,
     ):
         super().__init__()
         self._chunk_size = chunk_size
         self._continue_after_done = continue_after_done
+        self._allow_partial = allow_partial
 
     def is_ready(self, buffer_len) -> bool:
         return buffer_len >= self._chunk_size
@@ -176,6 +187,9 @@ def continues_in(self, consumer_walk: str | None) -> bool:
             return consumer_walk in self._continue_after_done
         return self._continue_after_done
 
+    def allow_partial_input(self) -> bool:
+        return self._allow_partial
+
 
 class RampChunkPolicy(ChunkPolicy):
     """A smaller first chunk, then chunks that grow geometrically up to a cap.

mstar/streaming/stream_buffer.py

@@ -168,10 +168,10 @@ def _join(self, pieces: list) -> torch.Tensor:
         pieces, concatenating and storing the result under a fresh uuid.
 
         A consumer that can take the pieces as they come should do that
-        instead, and set no ``join_dim`` -- which is where Qwen3-Omni's Talker
-        is headed. This is for a consumer that genuinely wants the whole
-        thing: it conditions on the complete prompt, or its shapes are
-        captured, and it has no use for a prefix.
+        instead (``ChunkingPolicy.per_chunk_outputs``), as Qwen3-Omni's Talker
+        does. This is for a consumer that genuinely wants the whole thing: it
+        conditions on the complete prompt, or its shapes are captured, and it
+        has no use for a prefix.
         """
         if len(pieces) == 1:
             return pieces[0]

mstar/utils/ipc_format.py

@@ -200,6 +200,12 @@ class ScheduleTPNode(MessageBody):
     # A row of the step takes an input a follower reads from elsewhere: every
     # rank agrees on whose read failed before preparing the step
     read_check: bool = False
+    # Per request_ids entry, the chunk the leader cut for the row: [start, end)
+    # of a measured total, -1 when the row is not chunked; all empty when no
+    # row is. A follower runs these rather than cutting its own.
+    chunk_starts: list[int] = field(default_factory=list)
+    chunk_ends: list[int] = field(default_factory=list)
+    chunk_totals: list[int] = field(default_factory=list)
 
 
 @dataclass

mstar/validation/placement.py

@@ -68,6 +68,7 @@ def check_placement(model: Model, config: Mapping[str, Any]) -> list[str]:
         *_walk_partitions(model, walk_graphs),
         *_node_placement(walk_graphs, node_groups),
         *_chunkable_nodes_are_looped(model, walk_graphs),
+        *_chunked_streams(model, walk_graphs),
     ]
 
 
@@ -119,6 +120,63 @@ def _chunkable_nodes_are_looped(
             )
 
 
+def _chunked_streams(
+    model: Model, walk_graphs: Mapping[str, GraphSection],
+) -> Iterator[str]:
+    """A chunked node's stream that its consumer would misread.
+
+    A ``per_chunk_outputs`` edge sends each chunk's piece as an item of its
+    own, so its consumer has to allow partial input. A gathered stream sends
+    every piece at once, which the consumer reads as one item only if the
+    connection joins them (``join_dim``) or it takes pieces anyway. A
+    consumer's ChunkPolicy counts items, so either mistake hands it the first
+    chunk as the whole input rather than failing.
+    """
+    from mstar.model.submodule_base import chunk_loop_name
+
+    connections = {
+        (conn.edge_name, conn.to_partition): conn
+        for conn in model.get_partition_topology().connections
+    }
+    for (node, walk), policy in model.chunking_policy_of().items():
+        if not policy.chunkable:
+            continue
+        loop = walk_graphs[walk].get_loops().get(chunk_loop_name(node, walk))
+        if loop is None:
+            continue  # _chunkable_nodes_are_looped reports it
+        where = f"node {node!r} in walk {walk!r}"
+        per_chunk = {e.name: e for e in loop.section.outputs if e.name in policy.per_chunk_outputs}
+        for name in sorted(policy.per_chunk_outputs - per_chunk.keys()):
+            yield f"{where} names per-chunk output {name!r}, which is not one of its outputs"
+        for edge in per_chunk.values():
+            if not edge.is_streaming:
+                yield (
+                    f"{where} sends {edge.name!r} per chunk, but it is not a streaming edge; "
+                    "only a streaming consumer can take a prompt's chunks one at a time"
+                )
+                continue
+            conn = connections.get((edge.name, getattr(edge, "target_partition", "")))
+            if conn is not None and not conn.chunk_policy_factory().allow_partial_input():
+                yield (
+                    f"{where} sends {edge.name!r} per chunk, but its consumer in partition "
+                    f"{conn.to_partition!r} does not allow partial input; set it "
+                    "(FixedChunkPolicy(allow_partial=True)), or gather the edge "
+                    "(accumulated_outputs) and join it (Connection.join_dim)"
+                )
+        for edge in loop.accumulated_outputs:
+            if not edge.is_streaming:
+                continue
+            conn = connections.get((edge.name, getattr(edge, "target_partition", "")))
+            if conn is not None and conn.join_dim is None \
+                    and not conn.chunk_policy_factory().allow_partial_input():
+                yield (
+                    f"{where} gathers stream {edge.name!r}, whose consumer in partition "
+                    f"{conn.to_partition!r} would read each chunk's piece as an item of "
+                    "its own; join them (Connection.join_dim), or send it per chunk "
+                    "(per_chunk_outputs) into a consumer that allows partial input"
+                )
+
+
 def _describe(wg: WorkerGraph) -> str:
     return (
         f"the worker graph of {sorted(wg.graph_walks)} with nodes "