Host overhead 3/N: worker speculates as columns, fused route+send, skips unread publishes; speculation and stream-buffer objects

configs/qwen3_5_0.8b.yaml

@@ -1,5 +1,10 @@
 model: "qwen3_5_0.8b"
 max_seq_len: 32768
+# The settings this model's throughput was measured with (mstar/utils/knobs.py)
+knobs:
+  kv_plan_cache: 1
+  ingraph_decode_rope: true
+  gc_gen0_threshold: 50000
 node_groups:
   # Text, images and video. Colocated: the
   # encoder runs once per request, then idles.

mstar/conductor/conductor.py

@@ -256,6 +256,7 @@ def _worker_process_target(
     tcp_transfer_device="",
     status_interval_s: float | None = None,
     env: dict[str, str] | None = None,
+    all_worker_graph_ids_to_instances: dict[int, list[list[str]]] | None = None,
 ):
     """Top-level target for spawned worker processes. Must be module-level for
     picklability. ``status_interval_s`` turns on the worker's metrics."""
@@ -301,6 +302,7 @@ def _graceful_exit(_signum, _frame):
             all_worker_graph_ids_to_graph_walks=all_worker_graph_ids_to_graph_walks,
             all_worker_graph_ids_to_nodes=all_worker_graph_ids_to_nodes,
             all_worker_graph_ids_to_dyn_loops=all_worker_graph_ids_to_dyn_loops,
+            all_worker_graph_ids_to_instances=all_worker_graph_ids_to_instances,
             sharding_config=sharding_config,
             parallel_groups=parallel_groups,
             hostname=hostname,
@@ -758,6 +760,15 @@ def _derive_worker_info(self):
             worker_graph_id: set(worker_graph.section.get_loops())
             for worker_graph_id, worker_graph in self.worker_graphs.items()
         }
+        # each worker graph's lockstep instances, as worker ids; a request of
+        # the worker graph runs on one of them
+        self._all_worker_graph_ids_to_instances: dict[int, list[list[str]]] = {
+            worker_graph_id: [
+                [f"worker_{rank}" for rank in instance]
+                for instance in (worker_graph._instance_ranks or [[r] for r in worker_graph.ranks])
+            ]
+            for worker_graph_id, worker_graph in self.worker_graphs.items()
+        }
         self._loop_max_iters: dict[str, int] = {}
         for worker_graph in self.worker_graphs.values():
             for name, loop in worker_graph.section.get_loops().items():
@@ -816,6 +827,7 @@ def _launch_workers(self):
                     "all_worker_graph_ids_to_graph_walks": self._all_worker_graph_ids_to_graph_walks,
                     "all_worker_graph_ids_to_nodes": self._all_worker_graph_ids_to_nodes,
                     "all_worker_graph_ids_to_dyn_loops": self._all_worker_graph_ids_to_dyn_loops,
+                    "all_worker_graph_ids_to_instances": self._all_worker_graph_ids_to_instances,
                     "sharding_config": self.per_worker_sharding_config[worker_id],
                     "parallel_groups": self.parallel_config.per_worker_config[worker_id],
                     "hostname": self.hostname,

mstar/engine/engine.py

@@ -6,6 +6,7 @@
 import time
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, field
+from itertools import repeat
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, NamedTuple
 
@@ -377,15 +378,14 @@ class ExecutingBatch:
     exec_timings: ExecTimings = field(default_factory=ExecTimings)
 
     def __post_init__(self):
-        self.per_request_info_wrapped = {
-            rid: ForwardPassInfoWrapper(
-                # Sparse: only rids with stream chunks or loop counters get an
-                # entry, so default like the .get() at prepare_inputs does.
-                info=info, metadata=self.per_request_input_metadata.get(
-                    rid, EMPTY_INPUT_METADATA,
-                ),
-            ) for rid, info in self.per_request_info.items()
-        }
+        # Sparse metadata: only rids with stream chunks or loop counters get
+        # an entry, so default like the .get() at prepare_inputs does. Built
+        # with C-level maps: this runs for every step, on every row.
+        rids = self.per_request_info.keys()
+        self.per_request_info_wrapped = dict(zip(rids, map(
+            ForwardPassInfoWrapper, self.per_request_info.values(),
+            map(self.per_request_input_metadata.get, rids, repeat(EMPTY_INPUT_METADATA)),
+        ), strict=True))
 
     @property
     def request_ids(self):
@@ -523,6 +523,8 @@ def load_model(
 
         node_names = set(submodules.keys())
         node_to_resources = {}
+        # resource -> every node declaring it, on this worker or another
+        declared_by: dict[str, frozenset[str]] = {}
         specs_by_key = resolve_spec_dependencies(specs)
 
         self._classify_tp_roles(node_names, parallel_groups)
@@ -564,6 +566,7 @@ def load_model(
             for node in relevant_nodes:
                 node_to_resources.setdefault(node, []).append(spec.resource_key)
             self._resource_nodes[spec.resource_key] = frozenset(relevant_nodes)
+            declared_by[spec.resource_key] = frozenset(spec.nodes)
 
         # rid -> (node, walk) of its last step here when that step was rows-only,
         # i.e. the sampler's slot master holds the token its next step reads
@@ -574,6 +577,7 @@ def load_model(
             # absent node would fall back to the full sweep
             node_resources={n: node_to_resources.get(n, []) for n in node_names},
             enable_nvtx=self._enable_nvtx,
+            declared_by=declared_by,
         )
 
         self._open_prefix_caches(specs_by_key, model)
@@ -2787,6 +2791,19 @@ def add_request(
             self.remove_request(request_id)
             raise
 
+    @property
+    def publisher_nodes(self) -> Mapping[str, frozenset[str]]:
+        """Each publishing resource's declaring nodes, across the deployment."""
+        return self._runner.publisher_nodes
+
+    def skip_publish(self, request_id, keys) -> None:
+        """Never publish ``keys`` for this request (see ``StepRunner.skip_publish``)."""
+        self._runner.skip_publish(request_id, keys)
+
+    def never_publish(self, keys) -> None:
+        """Never publish ``keys`` for any request (see ``StepRunner.never_publish``)."""
+        self._runner.never_publish(keys)
+
     def remove_request(self, request_id: str) -> set[str]:
         """Tear the request down on every resource. Returns the resources it
         held pages of that the teardown hands back, which `reclaimable` counts:

mstar/engine/resources/runner.py

@@ -65,6 +65,7 @@ def __init__(
         self, resources: Mapping[str, Resource],
         node_resources: Mapping[str, Collection[str]] | None = None,
         enable_nvtx: bool = False,
+        declared_by: Mapping[str, Collection[str]] | None = None,
     ):
         # Per-resource NVTX. `engine.plan.promoted` only reports that the sweep
         # was slow; these say WHICH resource ate it. Gated because range_push
@@ -114,6 +115,21 @@ def __init__(
         ]
         self._node_cap_order = self._per_node(node_resources, self._cap_order)
         self._node_trail_order = self._per_node(node_resources, self._trail_order)
+        # publisher key -> every node of the deployment that declares it
+        # (``declared_by``): a replica on another worker, under another node
+        # name, reads what this one publishes. This runner's own nodes when
+        # not given.
+        declared_by = declared_by or {}
+        self._publisher_nodes: dict[str, frozenset[str]] = {
+            key: frozenset(declared_by[key]) if key in declared_by else frozenset(
+                node for node, keys in (node_resources or {}).items() if key in keys
+            )
+            for key in self._publish_order
+        }
+        # publisher key -> the rids not to publish it for: those whose every
+        # declaring node runs on this lockstep instance, so no other instance
+        # can retrieve their state (``skip_publish``)
+        self._unpublished: dict[str, set] = {key: set() for key in self._publish_order}
         # capture-time buffer allocation, likewise scoped: a node's runner has
         # no business sizing a resource it never plans against
         self._node_order = self._per_node(node_resources, list(self._order))
@@ -349,6 +365,38 @@ def remove_request(self, rid: str) -> None:
         self._admission.remove_request(rid)
         for key in self._order:
             self._resources[key].remove_request(rid)
+        for rids in self._unpublished.values():
+            rids.discard(rid)
+
+    @property
+    def publisher_nodes(self) -> Mapping[str, frozenset[str]]:
+        """Each publishing resource's declaring nodes, across the deployment."""
+        return self._publisher_nodes
+
+    def skip_publish(self, rid, keys) -> None:
+        """Never publish ``keys`` for ``rid``: no other lockstep instance runs
+        a node that declares them for this request. Purged on removal."""
+        for key in keys:
+            self._unpublished[key].add(rid)
+
+    def never_publish(self, keys) -> None:
+        """``skip_publish`` for every request: ``keys`` leave the publish
+        order, so a step pays nothing for them."""
+        keys = set(keys)
+        self._publish_order = [k for k in self._publish_order if k not in keys]
+        if self._node_publish_order is not None:
+            self._node_publish_order = {
+                node: [k for k in order if k not in keys]
+                for node, order in self._node_publish_order.items()
+            }
+        for key in keys:
+            self._publisher_nodes.pop(key, None)
+            self._unpublished.pop(key, None)
+
+    def _publishing(self, order: list[str], request_ids) -> list[str]:
+        """``order`` minus the resources none of ``request_ids`` publishes."""
+        skip = self._unpublished
+        return [key for key in order if not (skip[key] and skip[key].issuperset(request_ids))]
 
     def admit_retrieve(
         self, rid: str, node_name: str, graph_walk: str,
@@ -610,14 +658,19 @@ def publish(
 
         Scoped to the node whose step just ran; no other node's state moved,
         and `merge_publish_info` keeps the entry it already published."""
-        order = self._sweep(self._node_publish_order, self._publish_order, node_name)
+        order = self._publishing(
+            self._sweep(self._node_publish_order, self._publish_order, node_name),
+            request_ids,
+        )
         if not order:
             return {rid: {} for rid in request_ids}
-        publishers = [(key, self._resources[key]) for key in order]
+        publishers = [(key, self._resources[key], self._unpublished[key]) for key in order]
         out: dict[str, dict[str, PublishedInfo]] = {}
         for rid in request_ids:
             per_key: dict[str, PublishedInfo] = {}
-            for key, resource in publishers:
+            for key, resource, skip in publishers:
+                if rid in skip:
+                    continue
                 info = resource.publish_for_step(
                     rid, node_name=node_name,
                     graph_walk=request_walks.get(rid, graph_walk) if request_walks else graph_walk,
@@ -640,7 +693,10 @@ def publish_snapshot(
         with nothing to publish has no entry. Under a combined walk each rid
         is snapshotted under its own walk (``request_walks``), as ``publish``
         does: a publisher may key what it publishes on the walk."""
-        order = self._sweep(self._node_publish_order, self._publish_order, node_name)
+        order = self._publishing(
+            self._sweep(self._node_publish_order, self._publish_order, node_name),
+            request_ids,
+        )
         if not order:
             return {}
         if request_walks:
@@ -663,7 +719,7 @@ def publish_snapshot(
                 per_key: dict[str, Any] = {}
                 for key, snaps in columns:
                     snap = snaps[i]
-                    if snap is not None:
+                    if snap is not None and rid not in self._unpublished[key]:
                         per_key[key] = snap
                 if per_key:
                     out[rid] = per_key
@@ -703,14 +759,19 @@ def publish_after_stop(
         request_walks: Mapping[int, str] | None = None,
     ) -> dict[str, dict[str, PublishedInfo]]:
         """Publish resources configured for the end of a dynamic loop."""
-        order = self._sweep(self._node_publish_order, self._publish_order, node_name)
+        order = self._publishing(
+            self._sweep(self._node_publish_order, self._publish_order, node_name),
+            request_ids,
+        )
         if not order:
             return {rid: {} for rid in request_ids}
-        publishers = [(key, self._resources[key]) for key in order]
+        publishers = [(key, self._resources[key], self._unpublished[key]) for key in order]
         out: dict[str, dict[str, PublishedInfo]] = {}
         for rid in request_ids:
             per_key: dict[str, PublishedInfo] = {}
-            for key, resource in publishers:
+            for key, resource, skip in publishers:
+                if rid in skip:
+                    continue
                 info = resource.publish_after_stop(
                     rid, node_name=node_name,
                     graph_walk=request_walks.get(rid, graph_walk) if request_walks else graph_walk,

mstar/graph/runtime/base.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from dataclasses import dataclass, field
+from enum import IntEnum
 from typing import NamedTuple
 
 from mstar.communication.tensors import NameToTensorList, TensorStore
@@ -411,6 +412,40 @@ class SpeculationPrepOutput(NamedTuple):
     spec_id: int = 0
 
 
+class SpecRows(NamedTuple):
+    """``GraphRuntime.prep_spec_rows``: a speculation's ready rows as columns."""
+    ready_rids: list[int]
+    # the worker graph every ready rid speculates in; -1 when none is ready
+    wg_id: int
+    # as SpeculationPrepOutput.spec_id; 0 here, since nothing streams in
+    spec_id: int
+    # None when no ready rid has a prepped input edge: everything the step
+    # reads comes from the in-flight step's outputs
+    input_edges: ColumnarEdgeSpecs | None
+
+
+class InlineRouteStatus(IntEnum):
+    """What ``route_and_send_inline`` did."""
+    # nothing (a rid with no host row, an unsupported dtype): the caller takes
+    # the general path
+    NOT_ROUTED = 0
+    # routed and sent
+    SENT = 1
+    # routed, but the send needs Python-side payloads: the caller finishes
+    # with ``send_outputs``
+    ROUTED = 2
+
+
+class InlineRoute(NamedTuple):
+    """``route_and_send_inline``'s result."""
+    status: InlineRouteStatus
+    completion_id: int
+    # ``send_outputs``' return, under SENT
+    sent_rids: list[int]
+    # the route for ``send_outputs`` to finish, under ROUTED
+    route_output: "RouteOutput | None"
+
+
 class ReadyNodeSpec(NamedTuple):
     node_name: str
     graph_walk: str
@@ -839,6 +874,34 @@ def prep_spec_rids(
         """
         pass
 
+    def prep_spec_rows(
+        self, spec_node_name: str, curr_node_name: str, graph_walk: str,
+        rids: list[int], room_for_continuing: int | None = None,
+    ) -> "SpecRows":
+        """``prep_spec_rids`` with no stream chunks to ingest, as columns, in
+        a form a runtime can serve in one crossing (the worker uses it where
+        ``SpeculationPreparer.rows_eligible`` allows). This default is that
+        composition over ``prep_spec_rids``; a runtime with a cheaper form
+        overrides it and must return the same.
+        """
+        prep = self.prep_spec_rids(SpeculationPrepInput(
+            spec_node_name=spec_node_name, curr_node_name=curr_node_name,
+            graph_walk=graph_walk, rids=list(rids),
+            room_for_continuing=room_for_continuing,
+            streaming_edges=[], streaming_edges_per_rid=[0] * len(rids),
+        ))
+        wg_ids = set(prep.wg_ids)
+        assert len(wg_ids) <= 1, (
+            "a speculation target lives in one worker graph, but its rows "
+            f"report several: {sorted(wg_ids)}"
+        )
+        return SpecRows(
+            ready_rids=list(prep.ready_rids),
+            wg_id=wg_ids.pop() if wg_ids else -1,
+            spec_id=prep.spec_id,
+            input_edges=prep.input_edges if len(prep.input_edges) else None,
+        )
+
     @abstractmethod
     def get_spec_target(
         self, curr_node_name: str, spec_node_name: str,