A stop token sampled at prefill ends the reply, and max_output_tokens counts the prefill's token

docs/adding_models.rst

@@ -1144,7 +1144,16 @@ methods are:
 ``check_stop(...) -> set[str] | StopCheck`` (optional)
    Runs off the GPU thread, and may read tensor values. Return the names of the ``Loop``
    objects to stop, for example after seeing the EOS token. This is how a decode loop
-   terminates. To also keep some of this step's outputs from being routed, return
+   terminates. A prefill samples the first token before its decode loop exists, so a stop
+   there is returned as ``END_PARTITION`` (from ``mstar.model.submodule_base``) instead:
+   the worker reports the partition done with that pass, and the conductor ends it rather
+   than start decode. Return it only from a row outside any loop; from inside one the
+   request fails. ``TokenStepNode`` does this for its ``PREFILL_WALK``. It suits a partition
+   with nothing left to run once generation stops: a model that runs another walk after its
+   loop (an ``align`` walk, or image generation after thinking) or publishes final-only
+   state at the loop's stop needs its own transition instead. Keep ``check_stop_batched``,
+   if you write one, in agreement.
+   To also keep some of this step's outputs from being routed, return
    ``StopCheck(stops, withheld=(names...))``, for example a frame streamed alongside the
    EOS token: its consumer then never receives it, instead of finding and dropping it with
    a device read. ``check_stop_batched`` returns a ``BatchedStopCheck`` for the same.

mstar/conductor/conductor.py

@@ -1793,6 +1793,11 @@ def _process_worker_graphs_done(
             )
             return []
 
+        # Any report of the pass may carry it, not only the one that completes
+        # the pass: that can be another worker's (an encoder's, say).
+        if body.partition_done:
+            pstate.worker_reported_done = True
+
         # Persist signals: every rank contributes its shard (different uuid +
         # source_tp_rank); accumulate across ranks, do not dedup.
         if body.persist_signals:
@@ -1851,10 +1856,7 @@ def _process_worker_graphs_done(
 
         return done_partitions
 
-    def _process_done_forward(
-        self, request_id: str, partition_name: str,
-        partition_done_from_worker: bool = False,
-    ) -> bool:
+    def _process_done_forward(self, request_id: str, partition_name: str) -> bool:
         """Process a completed forward pass for a specific partition.
 
         Calls get_partition_forward_pass_args for all partitions uniformly.
@@ -1865,10 +1867,16 @@ def _process_done_forward(
         """
         request_data = self.requests[request_id]
         pstate = request_data.partition_states[partition_name]
+        # A self-triggering partition's next pass can report in the same cycle,
+        # and is queued again before this pass's state is reset. When the
+        # earlier call ended the partition, on the flag the later report
+        # latched, the later call has nothing left to advance.
+        if pstate.is_done:
+            return all(ps.is_done for ps in request_data.partition_states.values())
 
         incoming_connections = request_data.get_incoming_connections(partition_name)
         stream_fed = partition_name in self._stream_fed_partitions
-        if stream_fed and partition_done_from_worker:
+        if stream_fed and pstate.worker_reported_done:
             pstate.is_done = True
 
         prev_walk =  pstate.metadata.graph_walk
@@ -1881,17 +1889,15 @@ def _process_done_forward(
         pstate.metadata = fwd_args.full_metadata
         pstate.metadata.kwargs.update(fwd_args.step_metadata)
 
-        # The worker's signal is authoritative for a stream-terminated
-        # partition: the pass that consumed the stream's final chunk has run.
-        # Treat it as the partition's own `request_done` so its downstream
-        # connections get `producer_done` and the chain closes. A model that
-        # inferred the end from the connection counters instead
-        # (`consumed_count >= token_count`) could race: `consumed_count` counts
-        # chunks popped for execution, and every worker-graphs-done report from
-        # a worker carries every buffer's count, so a colocated consumer's report
-        # could finish this partition while its last step was still running,
-        # and that step's output arrived for a request already gone.
-        if stream_fed and partition_done_from_worker:
+        # The worker's signal is authoritative, whatever walk the model chose:
+        # the pass that consumed a stream's final chunk, or whose stop check
+        # returned END_PARTITION, has run. Treat it as the partition's own
+        # `request_done` so its downstream connections get `producer_done`. A
+        # stream-fed model that inferred the end from the connection counters
+        # instead (`consumed_count >= token_count`) could race: a colocated
+        # consumer's report carries every buffer's count, so it could finish
+        # this partition while its last step was still running.
+        if pstate.worker_reported_done:
             fwd_args.request_done = True
 
         # Check max output tokens for partitions that produce tokens
@@ -1943,6 +1949,7 @@ def _process_done_forward(
         pstate.completed_worker_graph_ids = set()
         pstate.current_worker_graph_ids = set()
         pstate.wg_rank_completions = {}
+        pstate.worker_reported_done = False
         pstate.fwd_pass_number += 1
 
         self._set_partition_worker_graph_ids(
@@ -2084,7 +2091,7 @@ def _wait_for_workers_ready(self) -> None:
         logger.info("Conductor: all %d worker(s) ready", len(self.worker_ids))
 
     def _route_message(
-        self, message, done_partition_forwards: list[tuple[str, str, bool]],
+        self, message, done_partition_forwards: list[tuple[str, str]],
     ) -> None:
         """Apply one message from the api server or a worker. A WORKER_GRAPHS_DONE
         that finishes a partition's pass is queued on ``done_partition_forwards``."""
@@ -2119,9 +2126,7 @@ def _route_message(
                 range_push("conductor._process_worker_graphs_done")
             done_parts = self._process_worker_graphs_done(message.body)
             for pname in done_parts:
-                done_partition_forwards.append(
-                    (rid, pname, message.body.partition_done)
-                )
+                done_partition_forwards.append((rid, pname))
             if self.enable_nvtx:
                 range_pop()
         else:
@@ -2216,7 +2221,7 @@ def run(self):
                 range_push("conductor.run_loop")
 
             try:
-                done_partition_forwards: list[tuple[str, str, bool]] = []
+                done_partition_forwards: list[tuple[str, str]] = []
 
                 startup_backlog = self._startup_message_backlog
                 self._startup_message_backlog = []
@@ -2233,14 +2238,11 @@ def run(self):
 
                 completed_requests = []
 
-                for request_id, partition_name, p_done in done_partition_forwards:
+                for request_id, partition_name in done_partition_forwards:
                     if request_id not in self.requests or request_id in self.draining:
                         continue  # already completed by another partition in this cycle
                     try:
-                        all_done = self._process_done_forward(
-                            request_id, partition_name,
-                            partition_done_from_worker=p_done,
-                        )
+                        all_done = self._process_done_forward(request_id, partition_name)
                     except Exception as exc:
                         self._fail_after_error(
                             request_id, f"advancing partition {partition_name}", exc,

mstar/conductor/request_info.py

@@ -130,6 +130,9 @@ class PartitionState:
     current_worker_graph_ids: set[int] = field(default_factory=set)
     # wg_id -> count of distinct TP ranks that have reported completion
     wg_rank_completions: dict[int, int] = field(default_factory=dict)
+    # a report since the last advance set `partition_done` (it can be the
+    # next pass's; see `Conductor._process_done_forward`)
+    worker_reported_done: bool = False
     num_output_tokens: int = 0
     # resource label -> PublishedInfo, accumulated across all TP ranks and
     # handed back out on the next forward

mstar/engine/engine.py

@@ -310,9 +310,9 @@ class ExecutingBatch:
     # rid -> this step's InputMetadata; rids without any are absent
     per_request_input_metadata: Mapping[str, InputMetadata] = field(default_factory=dict)
 
-    # rids for which this step ends every finite stream into the partition;
-    # this step reports the partition done
-    stream_partition_done_rids: set[int] = field(default_factory=set)
+    # rids for which this step ends every finite stream into the partition, or
+    # whose stop check returned END_PARTITION; this step reports the partition done
+    partition_done_rids: set[int] = field(default_factory=set)
 
     # Populated on batch preparation
     inputs: list[NodeInputs] | None = None
@@ -435,7 +435,7 @@ def drop_rids(self, rids: set[int]):
         # A row that does not run this step must not end its streams on it:
         # the worker would stop its loop and finalize a row it no longer has.
         self.final_stream_rids.difference_update(rids)
-        self.stream_partition_done_rids.difference_update(rids)
+        self.partition_done_rids.difference_update(rids)
         self.per_request_info = {
             rid: info for rid, info in self.per_request_info.items()
             if rid not in rids
@@ -1908,7 +1908,8 @@ def check_stop_for_batch(
         self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList],
         host_rows: HostRows | None = None,
     ) -> dict[str, set[str]]:
-        """Each rid's ``submodule.check_stop``, as rid -> loops that should stop.
+        """Each rid's ``submodule.check_stop``, as rid -> its stops (loop names,
+        or ``END_PARTITION``).
 
         Reads tensor values, so it belongs on the caller's slow-postprocess
         path rather than in ``exec``, where the ``.item()`` / ``.cpu()`` would

mstar/graph/runtime/base.py

@@ -730,8 +730,9 @@ def cleanup_consumed_inputs(
         pass
 
     @abstractmethod
-    def mark_stream_partition_done(self, rid: int, partition: str):
-        """The consuming pass saw the final streaming chunk. Reported on the
+    def mark_partition_done(self, rid: int, partition: str):
+        """The partition is done: the consuming pass saw the final streaming
+        chunk, or a stop check returned ``END_PARTITION``. Reported on the
         WORKER_GRAPHS_DONE that follows, unless the node was speculative."""
         pass
 

mstar/graph/runtime/python.py

@@ -62,7 +62,7 @@
 class GraphRuntimePartitionInfo:
     graph_walk: str
     graph_walk_worker_graph_ids: list[int] = field(default_factory=list) # for this worker
-    stream_partition_done: bool = False  # set True when last chunk pops with is_final
+    partition_done: bool = False  # see GraphRuntime.mark_partition_done
 
 @dataclass
 class _SpecRidPrep:
@@ -452,10 +452,10 @@ def cleanup_consumed_inputs(
         # the caller.
         return FreedTensors.none()
 
-    def mark_stream_partition_done(self, rid: int, partition: str):
+    def mark_partition_done(self, rid: int, partition: str):
         info = self._request_info.get(rid)
         if info is not None and partition in info.partition_info:
-            info.partition_info[partition].stream_partition_done = True
+            info.partition_info[partition].partition_done = True
 
     def get_consumed_edges(
         self, source_node: str, dest_node: str, graph_walk: str,
@@ -1598,7 +1598,7 @@ def send_outputs(
                     # A speculatively-scheduled node has not really finished
                     # the partition, so it must not report done.
                     partition_done=(
-                        part is not None and part.stream_partition_done
+                        part is not None and part.partition_done
                         and not speculative
                     ),
                     profiling=profiling.get(rid),

mstar/graph/runtime/rust.py

@@ -358,8 +358,8 @@ def is_in_flight(
     ) -> bool:
         return self._rust.is_in_flight(node, wg_id, rid)
 
-    def mark_stream_partition_done(self, rid: int, partition: str):
-        self._rust.mark_stream_partition_done(rid, partition)
+    def mark_partition_done(self, rid: int, partition: str):
+        self._rust.mark_partition_done(rid, partition)
 
     def get_worker_graph_id_for_node(self, node: str, graph_walk: str) -> int:
         return self._rust.get_worker_graph_id_for_node(node, graph_walk)

mstar/model/bagel/bagel_model.py

@@ -1329,6 +1329,16 @@ def _get_step_metadata(
             "image_preprocess": full_metadata.kwargs["image_preprocess"],
         }
 
+    @staticmethod
+    def _prefill_token_metadata(sample_prefill_token: bool, target_output: str) -> dict:
+        """Whether this pass samples the reply's first token, and whether a stop
+        there ends the request: only for a text reply, which has no walk after
+        decode."""
+        return {
+            "sample_prefill_token": sample_prefill_token,
+            "prefill_stop_ends": sample_prefill_token and target_output == "text",
+        }
+
     def _get_fwd_pass_inputs(
         self,
         metadata: CurrentForwardConductorMetadata,
@@ -1460,7 +1470,7 @@ def get_initial_forward_pass_args(
         step_metadata =  self._get_step_metadata(full_metadata)
 
         sample_prefill_token = (think_mode or target_output != "image") and len(schedule) == 1
-        step_metadata["sample_prefill_token"] = sample_prefill_token
+        step_metadata.update(self._prefill_token_metadata(sample_prefill_token, target_output))
 
         inputs = self._get_fwd_pass_inputs(
             full_metadata, input_signals
@@ -1542,7 +1552,9 @@ def get_partition_forward_pass_args(
             request_done = True
 
         step_metadata =  self._get_step_metadata(metadata)
-        step_metadata["sample_prefill_token"] = sample_prefill_token
+        step_metadata.update(
+            self._prefill_token_metadata(sample_prefill_token, metadata.kwargs["target_output"])
+        )
         inputs = self._get_fwd_pass_inputs(
             metadata, persist_signals
         )

mstar/model/bagel/submodules.py

@@ -38,6 +38,7 @@
 from mstar.model.bagel.config import MIXED_TEXT, BagelModelConfig
 from mstar.model.higgs_audio.config import SAMPLER
 from mstar.model.submodule_base import (
+    END_PARTITION,
     ARNodeInputs,
     ARNodeSubmodule,
     BatchedModelOutput,
@@ -1677,8 +1678,15 @@ def check_stop(
             return set()
         token = outputs["new_token"][0].item()
         ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos
-        if (not ignore_eos and self.eos_token_id == token) or \
-                (request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 2 >= request_info.max_tokens):
+        hit_eos = not ignore_eos and self.eos_token_id == token
+        if request_info.graph_walk != "decode":
+            # The last prefill's token comes before decode_loop exists. The
+            # conductor says whether a stop there ends the request: not for an
+            # earlier prefill's throwaway token, nor for a think-then-image
+            # request, whose image_gen only the conductor's transition can start.
+            ends = request_info.step_metadata.get("prefill_stop_ends", False)
+            return {END_PARTITION} if hit_eos and ends else set()
+        if hit_eos or request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 2 >= request_info.max_tokens:
             return {"decode_loop"}
         return set()
 

mstar/model/canary/canary_model.py

@@ -172,7 +172,8 @@ def get_max_output_tokens_limit(self) -> int:
         return self.get_max_output_tokens()
 
     def get_loop_caps(self, **model_kwargs) -> dict[str, int]:
-        return {"decode_loop": self.get_max_output_tokens(**model_kwargs)}
+        # the prefill emits the first of max_output_tokens, so the loop runs one fewer
+        return {"decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)}
 
     def get_graph_walk_graphs(self) -> dict[str, GraphSection]:
         prefill = Sequential([

mstar/model/canary_qwen/canary_qwen_model.py

@@ -182,7 +182,8 @@ def get_max_output_tokens_limit(self) -> int:
         return self.MAX_OUTPUT_TOKENS
 
     def get_loop_caps(self, **model_kwargs) -> dict[str, int]:
-        return {DECODE_LOOP: self.get_max_output_tokens(**model_kwargs)}
+        # the prefill emits the first of max_output_tokens, so the loop runs one fewer
+        return {DECODE_LOOP: max(1, self.get_max_output_tokens(**model_kwargs) - 1)}
 
     def get_graph_walk_graphs(self) -> dict[str, GraphSection]:
         prefill = Sequential([

mstar/model/command_a_plus/submodules.py

@@ -29,6 +29,7 @@
     CommandAPlusTextConfig,
 )
 from mstar.model.submodule_base import (
+    END_PARTITION,
     ARNodeInputs,
     ARNodeSubmodule,
     ModelInputsFromEngine,
@@ -88,7 +89,6 @@ def __init__(
         self.language_model = language_model
         self.config = config
         self.image_token_id = image_token_id
-        self._ended_at_prefill: set[str] = set()
 
     def get_accelerator_graph_configs(
         self, device: torch.device, tp_world_size: int = 1,
@@ -207,20 +207,13 @@ def postprocess(self, request_id, request_info, outputs, **kwargs):
                 outputs["decode_input"] = outputs["new_token"]
             return
         outputs["text_inputs"] = outputs["new_token"]
-        if request_id in self._ended_at_prefill:
-            # The reply ended on its first token; this step only stops the loop.
-            del outputs["new_token"]
 
     def check_stop(self, request_id, request_info, outputs) -> set[str]:
         ignore_eos = request_info.resource_configs[SAMPLER].ignore_eos
         if request_info.graph_walk in PREFILL_WALKS:
-            if (request_info.max_tokens > 1 and not ignore_eos
-                    and outputs["new_token"][0].item() == self.config.eos_token_id):
-                self._ended_at_prefill.add(request_id)
-            return set()
-        if request_id in self._ended_at_prefill:
-            self._ended_at_prefill.discard(request_id)
-            return {"decode_loop"}
+            # the reply's first token, before decode_loop exists
+            eos = not ignore_eos and outputs["new_token"][0].item() == self.config.eos_token_id
+            return {END_PARTITION} if eos else set()
         if "new_token" not in outputs:
             return set()
         token = outputs["new_token"][0].item()
@@ -231,6 +224,3 @@ def check_stop(self, request_id, request_info, outputs) -> set[str]:
             return {"decode_loop"}
         return set()
 
-    def cleanup_request(self, request_id: str):
-        self._ended_at_prefill.discard(request_id)
-        super().cleanup_request(request_id)