From 24300580954c0e03ed8bd4c2ca546b8b31ce9123 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 001/219] resources: publish snapshot hooks so a step can defer its publication --- mstar/engine/resources/base.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index d5700b6ec0..462b74627b 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -214,6 +214,36 @@ def publish_after_stop( del request_id, node_name, graph_walk return None + def publish_snapshot_for_step( + self, + request_id: str, + node_name: str | None, + graph_walk: str | None, + ) -> Any: + """A cheap record of what ``publish_for_step`` would export right now. + + The worker finishes it with ``publish_from_snapshot`` only for the + requests whose frame or completion carries the publication, which in a + steady decode loop is none of them. The default takes the eager + publication itself, so a resource is correct without overriding; one + whose publish builds something per request (a page list, a transfer + descriptor) overrides both to record a few scalars here instead. + """ + return self.publish_for_step( + request_id, node_name=node_name, graph_walk=graph_walk, + ) + + def publish_from_snapshot( + self, + request_id: str, + snapshot: Any, + node_name: str | None, + graph_walk: str | None, + ) -> "PublishedInfo | None": + """Finish a ``publish_snapshot_for_step`` record into the published info.""" + del request_id, node_name, graph_walk + return snapshot + def reset_request(self, rid: str, free: bool=False): """For clearing dummy RIDs during cuda graph capture""" return From 91cae5084ae6f096da0d005bee8e306b96b6180f Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 002/219] runner: publish_snapshot and publish_from_snapshots sweeps --- mstar/engine/resources/runner.py | 50 ++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index 505b53d4b6..097db9c7a6 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -470,6 +470,56 @@ def publish( out[rid] = per_key return out + def publish_snapshot( + self, + request_ids: list[str], + node_name: str | None = None, + graph_walk: str | None = None, + ) -> dict[str, dict[str, Any]]: + """``publish``, deferred: per rid the snapshot each publisher took. + + Finished by ``publish_from_snapshots`` for the rids that need it. A rid + with nothing to publish has no entry.""" + order = self._sweep(self._node_publish_order, self._publish_order, node_name) + if not order: + return {} + publishers = [(key, self._resources[key]) for key in order] + out: dict[str, dict[str, Any]] = {} + for rid in request_ids: + per_key: dict[str, Any] = {} + for key, resource in publishers: + snap = resource.publish_snapshot_for_step( + rid, node_name=node_name, graph_walk=graph_walk, + ) + if snap is not None: + per_key[key] = snap + if per_key: + out[rid] = per_key + return out + + def publish_from_snapshots( + self, + snapshots: dict[str, dict[str, Any]], + request_ids: list[str], + node_name: str | None = None, + graph_walk: str | None = None, + ) -> dict[str, dict[str, PublishedInfo]]: + """The published info of ``request_ids`` from their snapshots.""" + out: dict[str, dict[str, PublishedInfo]] = {} + for rid in request_ids: + snaps = snapshots.get(rid) + if not snaps: + continue + per_key: dict[str, PublishedInfo] = {} + for key, snap in snaps.items(): + info = self._resources[key].publish_from_snapshot( + rid, snap, node_name=node_name, graph_walk=graph_walk, + ) + if info is not None: + per_key[key] = info + out[rid] = per_key + return out + def publish_after_stop( self, request_ids: list[str], From 68d7c4d086c4f75dd7843b3429bf3ccaaaacb516 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 003/219] kv: snapshot the publication per step and finish it on demand --- mstar/engine/resources/kv/manager.py | 70 ++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index c35c41f904..c041782418 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -1772,6 +1772,76 @@ def publish_for_step( request_id, node_name=node_name, graph_walk=graph_walk, ) + def publish_snapshot_for_step( + self, + request_id: str, + node_name: str | None, + graph_walk: str | None, + ): + """Per label the length, page count and generation ``publish`` would + export now. The page list itself is copied only when the publication + is finished, which in a steady decode loop is never: this ran for every + request on every step and was a quarter of the gpu thread's host time. + """ + with self._lock: + streams = self._streams.get(request_id) + overrides = self._overrides.get(request_id) + if streams is None or overrides is None: + return None + labels = overrides.get_publish_labels( + node_name, graph_walk, list(streams), final=False, + ) + if not labels: + return None + snap = tuple( + (label, stream.stored_len, len(stream.page_indices), + stream.reset_generation) + for label in labels + if (stream := streams.get(label)) is not None + ) + return snap or None + + def publish_from_snapshot( + self, + request_id: str, + snapshot, + node_name: str | None, + graph_walk: str | None, + ): + """``publish`` as of the snapshot: the lengths are the snapshot's, so a + step that committed in between (the next one is already running when + the worker finishes a publication) does not leak into it.""" + del node_name, graph_walk + if not snapshot: + return None + with self._lock: + streams = self._streams.get(request_id) + if streams is None: + return None + seq_info = {} + for label, seq_len, num_pages, reset_generation in snapshot: + stream = streams.get(label) + if stream is None: + continue + pages = list(stream.page_indices[:num_pages]) + seq_info[label] = KVSequenceInfo( + seq_len=seq_len, + latest_kv_transfer_info=self._transfer.get_kv_transfer_info( + request_id=request_id, + label=label, + page_indices=pages, + seq_len=seq_len, + reset_generation=reset_generation, + ), + page_indices=pages, + reset_generation=reset_generation, + ) + if not seq_info: + return None + return PublishedKVInfo.build_for_rank( + rank=self._rank, world_size=self._world_size, seq_info=seq_info, + ) + def publish_after_stop( self, request_id: str, From 18d8e174cbe7e82a2d7cd8b98bb8fae5d4244662 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 004/219] rope: snapshot the publication per step and finish it on demand --- mstar/engine/resources/position/manager.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/mstar/engine/resources/position/manager.py b/mstar/engine/resources/position/manager.py index 727029d1df..5a8d1c936e 100644 --- a/mstar/engine/resources/position/manager.py +++ b/mstar/engine/resources/position/manager.py @@ -195,6 +195,20 @@ def publish(self, request_id: str) -> "PublishedPositionInfo | None": return None return PublishedPositionInfo(counters=dict(counters)) + def publish_snapshot_for_step( + self, request_id: str, node_name: str | None, graph_walk: str | None, + ): + del node_name, graph_walk + counters = self._counters.get(request_id) + return tuple(counters.items()) if counters else None + + def publish_from_snapshot( + self, request_id: str, snapshot, node_name: str | None, + graph_walk: str | None, + ): + del request_id, node_name, graph_walk + return PublishedPositionInfo(counters=dict(snapshot)) if snapshot else None + def admit_retrieve( self, rid: str, node_name: str, graph_walk: str, published: "PublishedPositionInfo | None", From e84b7c2fb5df10fc10d65a9b674b6e8303265871 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 005/219] engine: finalize_batch snapshots the publication, materialize_publish finishes it --- mstar/engine/engine.py | 53 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 7b34c719f3..91298f7aab 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -75,6 +75,13 @@ # launch queue and block a launch (machine/driver dependent). _ENGINE_STEP_SYNC = os.environ.get("MSTAR_ENGINE_STEP_SYNC", "0") == "1" +# Publish resource state lazily: finalize_batch records a cheap snapshot per +# request and the worker finishes it only for the requests whose frame or +# completion carries it. Eager (0) builds the full publication for every +# request on every step; in a decode loop nothing reads it until the loop +# stops, and it was a quarter of the gpu thread's host time at batch 32. +_LAZY_PUBLISH = os.environ.get("MSTAR_LAZY_PUBLISH", "1") == "1" + def checkpoint_identity(path: str | Path) -> bytes: @@ -274,6 +281,9 @@ class ExecutingBatch: resource_publish_info: dict[str, dict[str, PublishedInfo]] = field( default_factory=dict ) + # What finalize_batch snapshotted instead of publishing (MSTAR_LAZY_PUBLISH); + # ``Engine.materialize_publish`` turns entries into resource_publish_info. + publish_snapshots: dict[str, dict[str, Any]] | None = None # The next step reads N's outputs, and plans against N's committed state. # Two separate dependencies, so two events: whoever prepares N+1 can start @@ -1378,10 +1388,23 @@ def finalize_batch( if self._enable_nvtx: range_push("engine.finalize_batch") try: + rids = ( + batch.request_ids if publish_request_ids is None + else publish_request_ids + ) + if _LAZY_PUBLISH: + # Snapshot now, on the gpu thread, so the lengths are this + # step's; finish in ``materialize_publish`` for the rids whose + # frame or completion carries the publication. + batch.publish_snapshots = self._runner.publish_snapshot( + rids, + node_name=batch.node_name, + graph_walk=batch.step_context.graph_walk, + ) + return batch.resource_publish_info # Returns rid -> {resource label -> published info} published = self._runner.publish( - batch.request_ids if publish_request_ids is None - else publish_request_ids, + rids, node_name=batch.node_name, graph_walk=batch.step_context.graph_walk, ) @@ -1394,6 +1417,32 @@ def finalize_batch( if self._enable_nvtx: range_pop() + def materialize_publish( + self, batch: ExecutingBatch, request_ids: list[int], + ) -> None: + """Finish the publication ``finalize_batch`` snapshotted, for the rids + whose frame or completion carries it. Idempotent per rid.""" + snaps = batch.publish_snapshots + if not snaps: + return + todo = [ + rid for rid in request_ids + if rid in snaps and rid not in batch.resource_publish_info + ] + if not todo: + return + published = self._runner.publish_from_snapshots( + snaps, todo, + node_name=batch.node_name, + graph_walk=batch.step_context.graph_walk, + ) + for rid in todo: + per_key = published.get(rid) or {} + batch.resource_publish_info[rid] = per_key + info = batch.per_request_info.get(rid) + if info is not None and per_key: + info.update_publish_info(per_key) + def finalize_stopped_requests( self, batch: ExecutingBatch, From 99593f15a8044e14481d66d33198d03960af8bd7 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 006/219] worker: finish publications only for stopped rids and rids that send a frame --- mstar/worker/worker.py | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 0998f50c4c..3ce5b3f973 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2835,18 +2835,22 @@ def _pp_stage(name: str) -> None: # Ordinary publication precedes stop detection on the GPU thread. # Export final-only state now, after the last iteration committed and - # before routing reports the completed loop to the conductor. + # before routing reports the completed loop to the conductor. The + # ordinary publication of a stopping request is finished first, so the + # completion carries both, as it did when every step published. if stopped_rids: + engine.materialize_publish(batch_N.node_batch, stopped_rids) engine.finalize_stopped_requests(batch_N.node_batch, stopped_rids) # CurrentForwardPassInfo also contains publication inherited from peer # ranks. Buffer only what this worker produced for the conductor. - for rid in batch_N.node_batch.request_ids: - self.request_state.buffer_publish_info( - rid, - batch_N.partition, - batch_N.node_batch.resource_publish_info.get(rid, {}), - ) + buffered_publish: set[int] = set() + for rid, published in batch_N.node_batch.resource_publish_info.items(): + if published: + self.request_state.buffer_publish_info( + rid, batch_N.partition, published, + ) + buffered_publish.add(rid) _pp_stage("stop_loops") if self.enable_nvtx: @@ -2984,6 +2988,19 @@ def _pp_stage(name: str) -> None: send_rids if needs_info is None else [rid for rid in send_rids if rid in needs_info] ) + # A frame to a peer or the conductor carries this step's publication: + # finish the snapshot for those rids (the stopped ones already are). + if info_rids: + engine.materialize_publish(batch_N.node_batch, info_rids) + for rid in info_rids: + if rid in buffered_publish: + continue + published = batch_N.node_batch.resource_publish_info.get(rid) + if published: + self.request_state.buffer_publish_info( + rid, batch_N.partition, published, + ) + buffered_publish.add(rid) send_input = SendInput( completion_id=route_output.completion_id, per_request_info=ParallelList( From 258e27a1a81368adc79e0af255f13521c5f5a234 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:56:24 -0700 Subject: [PATCH 007/219] test: lazy publication --- test/modular/test_lazy_publish.py | 245 ++++++++++++++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 test/modular/test_lazy_publish.py diff --git a/test/modular/test_lazy_publish.py b/test/modular/test_lazy_publish.py new file mode 100644 index 0000000000..949d0147d3 --- /dev/null +++ b/test/modular/test_lazy_publish.py @@ -0,0 +1,245 @@ +"""Lazy publication (MSTAR_LAZY_PUBLISH): finalize_batch snapshots a few scalars +per request and the worker finishes the publication only for the requests whose +frame or completion carries it. The finished publication has to equal what the +eager publish exported at snapshot time, even when the next step committed in +between (the next step is already running when the worker gets to it).""" +from __future__ import annotations + +import pytest +import torch + +from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine import engine as engine_mod +from mstar.engine.engine import Engine, ExecutingBatch +from mstar.engine.resources import StepRunner +from mstar.engine.resources.base import Resource +from mstar.engine.resources.kv import manager as manager_mod +from mstar.engine.resources.kv.config import KVStep, PagedKVConfig +from mstar.engine.resources.kv.manager import KVManager +from mstar.engine.resources.position.manager import ( + PublishedPositionInfo, + RopeManager, +) +from mstar.engine.resources.step import ADMIT_OK, Segment, StepContext + +OWN_HANDLE = ("own", 0) + + +class _StubTransfer: + def __init__(self, transfer_engine_info, kv_cache, **kwargs): + del transfer_engine_info, kv_cache, kwargs + + def get_kv_transfer_info(self, **kwargs): + # what the real engines stamp: the lengths the caller handed over + return (OWN_HANDLE, kwargs.get("seq_len"), tuple(kwargs.get("page_indices") or ())) + + def start_async_retrieve(self, **kwargs): + del kwargs + + def owns_transfer_info(self, transfer_info, **kwargs): + del kwargs + return transfer_info[0] == OWN_HANDLE + + def cleanup(self): + pass + + def remove_request(self, request_id): + del request_id + + +@pytest.fixture(autouse=True) +def _stub(monkeypatch): + monkeypatch.setattr(manager_mod, "KVTransferManager", _StubTransfer) + + +def _manager(max_num_pages=64, page_size=16) -> KVManager: + return KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=4096, + max_num_pages=max_num_pages, page_size=page_size, + ), + name="kv", joint_comm_group=None, transfer_engine_info=None, + device=torch.device("cpu"), dtype=torch.float32, + ) + + +def _ctx(*rids): + return StepContext(request_ids=tuple(rids), graph_walk="w", slot=0, capture=False) + + +def _grow(kv, rid, span, label="main"): + step = KVStep(segments=(Segment(rid, label, span),)) + ctx = _ctx(rid) + assert kv.admit(step, ctx).ok + kv.commit(step, ctx) + + +def _seq(published, label="main"): + return published.get(0)[label] + + +# -- KV ---------------------------------------------------------------------- + + +def test_kv_snapshot_finishes_to_the_eager_publication(): + kv = _manager() + kv.ingest_request("r0", None) + _grow(kv, "r0", 40) + eager = kv.publish("r0", "node", "w") + snap = kv.publish_snapshot_for_step("r0", "node", "w") + lazy = kv.publish_from_snapshot("r0", snap, "node", "w") + assert _seq(lazy).seq_len == _seq(eager).seq_len == 40 + assert _seq(lazy).page_indices == _seq(eager).page_indices + assert _seq(lazy).reset_generation == _seq(eager).reset_generation + assert _seq(lazy).latest_kv_transfer_info == _seq(eager).latest_kv_transfer_info + assert lazy.world_size == eager.world_size + + +def test_kv_snapshot_keeps_this_steps_lengths_when_the_next_step_committed(): + kv = _manager(page_size=16) + kv.ingest_request("r0", None) + _grow(kv, "r0", 31) + expected = kv.publish("r0", "node", "w") + snap = kv.publish_snapshot_for_step("r0", "node", "w") + # the next step lands (crosses a page boundary: a third page appears) + _grow(kv, "r0", 1) + _grow(kv, "r0", 1) + later = kv.publish("r0", "node", "w") + assert _seq(later).seq_len == 33 and len(_seq(later).page_indices) == 3 + lazy = kv.publish_from_snapshot("r0", snap, "node", "w") + assert _seq(lazy).seq_len == _seq(expected).seq_len == 31 + assert _seq(lazy).page_indices == _seq(expected).page_indices + assert len(_seq(lazy).page_indices) == 2 + # the transfer descriptor is stamped with the snapshot's lengths too + assert _seq(lazy).latest_kv_transfer_info == _seq(expected).latest_kv_transfer_info + + +def test_kv_snapshot_is_a_copy_not_a_view_of_the_page_list(): + kv = _manager() + kv.ingest_request("r0", None) + _grow(kv, "r0", 20) + snap = kv.publish_snapshot_for_step("r0", "node", "w") + lazy = kv.publish_from_snapshot("r0", snap, "node", "w") + pages = list(_seq(lazy).page_indices) + _grow(kv, "r0", 40) + assert _seq(lazy).page_indices == pages + + +def test_kv_snapshot_is_none_for_an_unknown_request_and_mirrors_an_empty_stream(): + kv = _manager() + assert kv.publish_snapshot_for_step("nobody", "node", "w") is None + assert kv.publish_from_snapshot("r0", None, "node", "w") is None + kv.ingest_request("r0", None) + # ingest opens the stream at length 0; eager publish exports that, and so does the snapshot + eager = kv.publish("r0", "node", "w") + snap = kv.publish_snapshot_for_step("r0", "node", "w") + lazy = kv.publish_from_snapshot("r0", snap, "node", "w") + assert _seq(lazy).seq_len == _seq(eager).seq_len == 0 + assert _seq(lazy).page_indices == _seq(eager).page_indices == [] + + +# -- positions ----------------------------------------------------------------- + + +def test_position_snapshot_matches_publish_and_does_not_track_later_advances(): + pm = RopeManager.__new__(RopeManager) + pm._counters = {"r0": {"main": 17}} + eager = pm.publish("r0") + snap = pm.publish_snapshot_for_step("r0", "node", "w") + pm._counters["r0"]["main"] = 18 + lazy = pm.publish_from_snapshot("r0", snap, "node", "w") + assert isinstance(lazy, PublishedPositionInfo) + assert lazy.counters == eager.counters == {"main": 17} + assert pm.publish_snapshot_for_step("r1", "node", "w") is None + assert pm.publish_from_snapshot("r1", None, "node", "w") is None + + +# -- runner + engine ----------------------------------------------------------- + + +class _Pub(Resource): + """A resource with only the default hooks: its snapshot is the eager value.""" + + def __init__(self, published): + self._published = published + self.publish_calls = 0 + + @classmethod + def build(cls, spec, info): + raise NotImplementedError + + def admit(self, step, ctx): + return ADMIT_OK + + def plan(self, step, ctx): + return None + + def publish(self, request_id): + self.publish_calls += 1 + return None if self._published is None else f"{self._published}:{request_id}" + + +def test_runner_snapshot_and_finish_equal_publish_and_skip_the_silent(): + kv, sampler = _Pub("kv-info"), _Pub(None) + runner = StepRunner({"kv": kv, "sampler": sampler}) + eager = runner.publish(["r1", "r2"]) + snaps = runner.publish_snapshot(["r1", "r2"]) + assert set(snaps) == {"r1", "r2"} and set(snaps["r1"]) == {"kv"} + assert runner.publish_from_snapshots(snaps, ["r2"]) == {"r2": eager["r2"]} + assert runner.publish_from_snapshots(snaps, ["r9"]) == {} + + +def _engine_with(runner): + eng = Engine.__new__(Engine) + eng._runner = runner + eng._enable_nvtx = False + return eng + + +def _batch(rids): + return ExecutingBatch( + node_name="LLM", + per_request_info={ + rid: CurrentForwardPassInfo( + request_id=str(rid), graph_walk="decode", fwd_index=0, + random_seed=0, max_tokens=16, + ) + for rid in rids + }, + step_context=StepContext( + request_ids=tuple(rids), graph_walk="decode", slot=0, capture=False, + ), + ) + + +def test_lazy_finalize_publishes_nothing_until_materialized(monkeypatch): + monkeypatch.setattr(engine_mod, "_LAZY_PUBLISH", True) + kv = _Pub("kv-info") + eng = _engine_with(StepRunner({"kv": kv})) + batch = _batch([1, 2, 3]) + out = eng.finalize_batch(batch) + assert out == {} and batch.resource_publish_info == {} + assert set(batch.publish_snapshots) == {1, 2, 3} + assert all(not info.resource_publish_info for info in batch.per_request_info.values()) + + eng.materialize_publish(batch, [2]) + assert batch.resource_publish_info == {2: {"kv": "kv-info:2"}} + assert batch.per_request_info[2].resource_publish_info == {"kv": "kv-info:2"} + assert not batch.per_request_info[1].resource_publish_info + # idempotent, and a rid the batch never ran is ignored + calls = kv.publish_calls + eng.materialize_publish(batch, [2, 7]) + assert kv.publish_calls == calls + assert batch.resource_publish_info == {2: {"kv": "kv-info:2"}} + + +def test_eager_finalize_still_publishes_every_request(monkeypatch): + monkeypatch.setattr(engine_mod, "_LAZY_PUBLISH", False) + kv = _Pub("kv-info") + eng = _engine_with(StepRunner({"kv": kv})) + batch = _batch([1, 2]) + out = eng.finalize_batch(batch) + assert out == {1: {"kv": "kv-info:1"}, 2: {"kv": "kv-info:2"}} + assert batch.per_request_info[1].resource_publish_info == {"kv": "kv-info:1"} + assert batch.publish_snapshots is None + eng.materialize_publish(batch, [1]) # nothing to finish, no error From b36967b0ed4e71a6216517919843a0a5a6ac00fe Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:58:09 -0700 Subject: [PATCH 008/219] submodule_base: a stop buffer that is a row output shares the row clone --- mstar/model/submodule_base.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 45c964f968..a010dae8b4 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -109,14 +109,29 @@ def pop(self, request_id: str, default=None): alone, and the caller takes the request out of the batch instead.""" return self.per_rid_outputs.pop(request_id, default) - def clone_check_stop_buffers(self): + def clone_check_stop_buffers( + self, reuse: dict[str, torch.Tensor] | None = None, + ): """A detached copy: the next replay overwrites a captured graph's - buffers before the stop check reads them.""" + buffers before the stop check reads them. + + ``reuse`` holds clones already taken of the row outputs (see + ``clone_row_outputs``); a stop buffer that is the very same tensor as a + row output takes that clone instead of a second copy. The row clone is + narrowed to the real rows, which is every row the stop check reads. + """ if self.check_stop_buffers is None: return None + rows = self.row_outputs or {} - def _clone(value): + def _clone(value, name=None): if isinstance(value, torch.Tensor): + if ( + reuse is not None and name is not None + and value is rows.get(name) + and isinstance(reuse.get(name), torch.Tensor) + ): + return reuse[name] return value.clone() if isinstance(value, list): return [ @@ -127,7 +142,7 @@ def _clone(value): # pass anything else through: dropping it could lose a stop signal return value - return {k: _clone(v) for k, v in self.check_stop_buffers.items()} + return {k: _clone(v, k) for k, v in self.check_stop_buffers.items()} def clone_row_outputs(self, num_rows: int) -> dict[str, torch.Tensor]: """One clone per row-addressed batch tensor, narrowed to the real rows From 52b171e7eef716c205438e1ee5088d23e27ef2e0 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:58:09 -0700 Subject: [PATCH 009/219] engine: collect_outputs fast path for row-only outputs --- mstar/engine/engine.py | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 91298f7aab..b5f1af2314 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1486,7 +1486,7 @@ def _collect_outputs( ) outputs: dict[str, NameToTensorList] = {} - self._merge_per_rid( + row_clones = self._merge_per_rid( outputs, raw_outputs, request_ids, out_ids, submodule, req_info, ) self._merge_unpacked( @@ -1495,7 +1495,11 @@ def _collect_outputs( ) return BatchedModelOutput( per_rid_outputs=outputs, - check_stop_buffers=raw_outputs.clone_check_stop_buffers(), + # the stop buffer that is a row output shares its clone: one copy + # of the sampled tokens per step, not two + check_stop_buffers=raw_outputs.clone_check_stop_buffers( + reuse=row_clones, + ), # row i is request_ids[i] as the forward ran them, not the # worker's later-rewritten request list row_request_ids=tuple(request_ids), @@ -1509,14 +1513,37 @@ def _merge_per_rid( out_ids: list[str], submodule: NodeSubmodule, req_info: Mapping[str, CurrentForwardPassInfo], - ) -> None: + ) -> dict[str, torch.Tensor]: """Fold the forward's per-rid entries into ``outputs``. ``row_outputs`` are cloned once for the batch and handed out as row views; per-rid entries are cloned individually after ``filter_batched_output``. A name in both resolves to the per-rid one. + Returns the row clones so the caller can share them. """ row_clones = raw_outputs.clone_row_outputs(len(request_ids)) + # The common decode case: only row outputs, and a submodule that does + # not filter per request. One ``split`` per output makes every row + # view in a single call, the same ``tensor[i:i+1]`` views the loop + # below would slice one at a time, and the per-rid dicts are built + # without the candidate/from_rows bookkeeping. + if ( + not raw_outputs.per_rid_outputs + and type(submodule).filter_batched_output + is NodeSubmodule.filter_batched_output + ): + row_views = { + name: tensor.split(1) + for name, tensor in row_clones.items() + if isinstance(tensor, torch.Tensor) and tensor.dim() + } + if row_views: + for i, rid in enumerate(request_ids): + merged = outputs.setdefault(rid, {}) + for name, views in row_views.items(): + if i < len(views): + merged[name] = [views[i]] + return row_clones for i, (rid, out_id) in enumerate( zip(request_ids, out_ids, strict=False) ): @@ -1549,6 +1576,7 @@ def _merge_per_rid( merged[key] = [value.clone()] else: merged[key] = value + return row_clones def _merge_unpacked( self, From 927224610065631fa915539bfdecab50735592c9 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 15:58:09 -0700 Subject: [PATCH 010/219] test: collect_outputs fast path --- .../modular/test_collect_outputs_fast_path.py | 120 ++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 test/modular/test_collect_outputs_fast_path.py diff --git a/test/modular/test_collect_outputs_fast_path.py b/test/modular/test_collect_outputs_fast_path.py new file mode 100644 index 0000000000..4db4db772d --- /dev/null +++ b/test/modular/test_collect_outputs_fast_path.py @@ -0,0 +1,120 @@ +"""Engine._collect_outputs: the decode fast path (row outputs only, no per-rid +filtering) hands out the same per-rid row views as the generic walk, and the +stop buffer shares the row clone instead of taking a second copy.""" +from __future__ import annotations + +import torch + +from mstar.engine.engine import Engine +from mstar.model.submodule_base import BatchedModelOutput, NodeSubmodule + + +class _Sub(NodeSubmodule): + """The base filter: a no-op, which is what takes the fast path.""" + + def __init__(self): + pass + + def prepare_inputs(self, *args, **kwargs): + raise NotImplementedError + + def forward(self, *args, **kwargs): + raise NotImplementedError + + +class _Filtering(_Sub): + """Overrides the filter, so the generic per-rid walk has to run.""" + + def filter_batched_output(self, request_info, outputs): + return {k: v for k, v in outputs.items() if k != "aux"} + + +def _raw(bs=4, with_per_rid=False): + tok = torch.arange(bs, dtype=torch.int64) + 100 + aux = torch.arange(bs * 2, dtype=torch.float32).reshape(bs, 2) + raw = BatchedModelOutput( + row_outputs={"new_token": tok, "aux": aux}, + check_stop_buffers={"new_token": tok}, + ) + if with_per_rid: + raw.per_rid_outputs = {1: {"extra": [torch.tensor([7])]}} + return raw + + +def _collect(sub, raw, rids): + eng = Engine.__new__(Engine) + outputs = {} + row_clones = eng._merge_per_rid(outputs, raw, rids, rids, sub, {}) + return outputs, row_clones + + +def _generic_reference(raw, rids, sub): + """The per-rid walk as it was before the fast path.""" + row_clones = raw.clone_row_outputs(len(rids)) + outputs = {} + for i, rid in enumerate(rids): + candidates = { + name: [t[i:i + 1]] for name, t in row_clones.items() + if isinstance(t, torch.Tensor) and t.dim() and i < t.shape[0] + } + rid_out = sub.filter_batched_output(None, candidates) + outputs[rid] = {k: v for k, v in rid_out.items()} + return outputs + + +def _same(a, b): + assert a.keys() == b.keys() + for rid in a: + assert a[rid].keys() == b[rid].keys(), rid + for name in a[rid]: + va, vb = a[rid][name], b[rid][name] + assert len(va) == len(vb) == 1 + assert va[0].shape == vb[0].shape + assert torch.equal(va[0], vb[0]) + + +def test_fast_path_matches_the_generic_walk(): + rids = [10, 11, 12, 13] + raw = _raw() + fast, _ = _collect(_Sub(), raw, rids) + _same(fast, _generic_reference(raw, rids, _Sub())) + # row views, not copies: a row is a [1] view into the batch clone + assert fast[11]["new_token"][0].shape == (1,) and fast[11]["new_token"][0].item() == 101 + assert fast[12]["aux"][0].shape == (1, 2) + + +def test_fast_path_handles_fewer_rows_than_requests_and_padding(): + rids = [1, 2] + raw = _raw(bs=4) # a padded bucket: four rows, two real requests + fast, clones = _collect(_Sub(), raw, rids) + assert set(fast) == {1, 2} + assert clones["new_token"].shape == (2,), "narrowed to the real rows" + assert fast[2]["new_token"][0].item() == 101 + + +def test_per_rid_entries_or_a_filtering_submodule_take_the_generic_walk(): + rids = [1, 2, 3] + raw = _raw(bs=3, with_per_rid=True) + out, _ = _collect(_Sub(), raw, rids) + assert "extra" in out[1] and torch.equal(out[1]["extra"][0], torch.tensor([7])) + assert out[1]["extra"][0] is not raw.per_rid_outputs[1]["extra"][0], "per-rid entries are cloned" + filt, _ = _collect(_Filtering(), _raw(bs=3), rids) + assert all("aux" not in filt[r] and "new_token" in filt[r] for r in rids) + + +def test_stop_buffer_shares_the_row_clone(): + raw = _raw(bs=4) + clones = raw.clone_row_outputs(3) + stop = raw.clone_check_stop_buffers(reuse=clones) + assert stop["new_token"] is clones["new_token"] + assert stop["new_token"].shape == (3,) + # without reuse, or for a buffer that is not a row output, it is still a fresh copy + fresh = raw.clone_check_stop_buffers() + assert fresh["new_token"] is not raw.check_stop_buffers["new_token"] + assert torch.equal(fresh["new_token"], raw.check_stop_buffers["new_token"]) + other = BatchedModelOutput( + row_outputs={"new_token": torch.zeros(2, dtype=torch.int64)}, + check_stop_buffers={"new_token": torch.ones(2, dtype=torch.int64)}, + ) + stop2 = other.clone_check_stop_buffers(reuse=other.clone_row_outputs(2)) + assert torch.equal(stop2["new_token"], torch.ones(2, dtype=torch.int64)) From c05f8bb67c8f3ed4aa795038091f6200b707877a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:00:26 -0700 Subject: [PATCH 011/219] micro_scheduler: check backlog readiness only up to a step's worth --- mstar/worker/micro_scheduler.py | 39 +++++++++++++++++++++++++++------ 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/mstar/worker/micro_scheduler.py b/mstar/worker/micro_scheduler.py index 4788550eef..91f60e2a9e 100644 --- a/mstar/worker/micro_scheduler.py +++ b/mstar/worker/micro_scheduler.py @@ -605,12 +605,40 @@ def _filter_cap_and_schedule( capture_group_of: int | None = None, ): node_partition = request_state.get_partition_for_node(batch.node_name) - not_ready_rids = { - rid for rid in batch.request_to_worker_graph if not self._check_ready( + # Check readiness in backlog order and stop once a step's worth is + # found: `_cap_batch_and_schedule` takes the first `max_bs` rids not + # excluded, which are exactly the ones checked and found ready, and + # the rids past that point keep their place for the next pass. Before + # this every backlogged request was checked (an engine sweep over the + # resources each) on every step, with 96 waiting at 128 in flight. + limit = len(batch) if max_bs is None else max(max_bs, 0) + group = functools.partial( + self._capture_group, request_state, + batch.node_name, batch.graph_walk, + ) + # `is not None`: handle 0 is a real rid + wanted = None if capture_group_of is None else group(capture_group_of) + anchored = capture_group_of is not None + not_ready_rids: set[int] = set() + off_group: set[int] = set() + taken = 0 + for rid in batch.request_to_worker_graph: + if taken >= limit: + break + if not self._check_ready( batch.node_name, rid, request_state.get_fwd_info(rid, node_partition), - ) - } + ): + not_ready_rids.add(rid) + continue + rid_group = group(rid) + if not anchored: + # the first eligible rid sets the group, as before + wanted, anchored = rid_group, True + if rid_group != wanted: + off_group.add(rid) + continue + taken += 1 # A failed rid is not "not ready yet": excluding it would put it # straight back in the backlog. This chunk is out of `self.backlog` # right now, so `_drop_backlogged_rid` cannot reach it. @@ -622,9 +650,6 @@ def _filter_cap_and_schedule( batch.request_to_worker_graph.keys() ) not_ready_rids -= self.failed_rids - off_group = self._off_group_rids( - request_state, batch, not_ready_rids, anchor=capture_group_of, - ) return self._cap_batch_and_schedule( batch, max_bs, not_ready_rids | off_group, ) From 219103522d965c14d777dc02c84e5b51cbf4d660 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:00:27 -0700 Subject: [PATCH 012/219] test: bounded backlog readiness checks --- test/modular/test_micro_scheduler.py | 66 ++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/test/modular/test_micro_scheduler.py b/test/modular/test_micro_scheduler.py index aafe250d25..711a3e0f65 100644 --- a/test/modular/test_micro_scheduler.py +++ b/test/modular/test_micro_scheduler.py @@ -759,3 +759,69 @@ def test_a_submodule_that_does_not_opt_in_is_never_split(): def test_a_walk_without_a_capture_is_never_split(): assert _capture_group(captured=()) is None + + +# ── bounded readiness checks on the backlog ────────────────────────────── + + +class _CountingEngine(_Engine): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.checked: list[str] = [] + + def check_ready(self, node_name, rid, fwd_info, allow_reload=True): + self.checked.append(rid) + return super().check_ready(node_name, rid, fwd_info, allow_reload) + + +def test_the_backlog_is_checked_only_up_to_a_steps_worth(): + """A step takes the first `max_bs` ready rids in backlog order; the rids + past that point are not asked (an engine sweep each) and keep their place.""" + engine = _CountingEngine(max_bs=2) + sched = _scheduler(engine) + sched.backlog[(NODE, WALK)] = _batch([f"r{i}" for i in range(8)]) + + batch = _next_batch(sched, _Manager([])) + + assert list(batch.request_to_worker_graph) == ["r0", "r1"] + assert engine.checked == ["r0", "r1"] + assert list(sched.backlog[(NODE, WALK)].request_to_worker_graph) == [ + f"r{i}" for i in range(2, 8) + ] + + +def test_an_unready_rid_is_skipped_and_the_scan_goes_one_further(): + engine = _CountingEngine(max_bs=2, not_ready={"r1"}) + sched = _scheduler(engine) + sched.backlog[(NODE, WALK)] = _batch([f"r{i}" for i in range(6)]) + + batch = _next_batch(sched, _Manager([])) + + assert list(batch.request_to_worker_graph) == ["r0", "r2"] + assert engine.checked == ["r0", "r1", "r2"] + # the unready rid goes to the back, the unchecked ones keep their order + assert list(sched.backlog[(NODE, WALK)].request_to_worker_graph) == [ + "r3", "r4", "r5", "r1", + ] + + +def test_a_full_caller_batch_checks_nothing_in_the_backlog(): + engine = _CountingEngine(max_bs=2) + sched = _scheduler(engine) + sched.backlog[(NODE, WALK)] = _batch(["r0", "r1", "r2"]) + + assert _next_batch(sched, _Manager([]), pre_existing_batch_size=2) is None + assert engine.checked == [] + assert list(sched.backlog[(NODE, WALK)].request_to_worker_graph) == ["r0", "r1", "r2"] + + +def test_the_group_anchor_is_applied_to_the_checked_rids_only(): + engine = _CountingEngine(max_bs=2, groups={"r0": "a", "r1": "b", "r2": "a", "r3": "a"}) + sched = _scheduler(engine) + sched.backlog[(NODE, WALK)] = _batch(["r0", "r1", "r2", "r3"]) + + batch = _next_batch(sched, _Manager([])) + + assert list(batch.request_to_worker_graph) == ["r0", "r2"] + assert engine.checked == ["r0", "r1", "r2"] + assert list(sched.backlog[(NODE, WALK)].request_to_worker_graph) == ["r3", "r1"] From f5a732716627600af179f4dfa803be53f79b326a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:03:36 -0700 Subject: [PATCH 013/219] resources: keys_prefix_chains says whether a prefix cache is open --- mstar/engine/resources/base.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index 462b74627b..6ea4947740 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -128,6 +128,12 @@ def extend_prefix_chain( """Take this step's sampled tokens, so what was generated can be keyed.""" return + @property + def keys_prefix_chains(self) -> bool: + """Whether ``extend_prefix_chain`` can do anything right now (a prefix + cache is open). The runner skips the per-request sweep otherwise.""" + return False + def fingerprint(self) -> bytes | None: """What this resource contributes to the prefix cache's root. From 9e38073d02efc3a726059418fc109bef58d9c5ec Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:03:36 -0700 Subject: [PATCH 014/219] kv: report an open prefix cache through keys_prefix_chains --- mstar/engine/resources/kv/manager.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index c041782418..318629a68b 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -676,6 +676,10 @@ def _seed_keys(self, rid: str, label: str, stream: CacheStream) -> None: tail = list((overrides.prefix_tail or {}).get(label) or ()) stream.chain = PrefixChain.seed(keys, tail, self.config.page_size) + @property + def keys_prefix_chains(self) -> bool: + return self._index is not None + def extend_prefix_chain( self, rid: str, node_name: str, graph_walk: str, outputs, ) -> None: From f1b0be51af68a0b391e986b86a7805db81f24715 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:03:36 -0700 Subject: [PATCH 015/219] runner: prefix_chain_keys, and extend_prefix_chains over a given key list --- mstar/engine/resources/runner.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index 097db9c7a6..95f947cc02 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -135,11 +135,23 @@ def apply_cached_prefix( rid, node_name, graph_walk, inputs, matched_len, ) + def prefix_chain_keys(self, node_name: str | None) -> list[str]: + """The node's resources with a prefix cache open right now. Empty in + the common case, which lets the caller skip a per-request sweep that + was six no-op calls per request per step.""" + return [ + key for key in self._sweep(self._node_order, self._order, node_name) + if self._resources[key].keys_prefix_chains + ] + def extend_prefix_chains( self, rid: str, node_name: str, graph_walk: str, outputs, + keys: list[str] | None = None, ) -> None: """Offer this step's sampled tokens to the node's own resources.""" - for key in self._sweep(self._node_order, self._order, node_name): + if keys is None: + keys = self._sweep(self._node_order, self._order, node_name) + for key in keys: self._resources[key].extend_prefix_chain( rid, node_name, graph_walk, outputs, ) From 5ef8de3baafca56af77dfde64b9370d57769b9f5 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:03:36 -0700 Subject: [PATCH 016/219] engine: skip the per-request prefix-chain sweep when no cache is open --- mstar/engine/engine.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index b5f1af2314..c35d49cc9d 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -755,12 +755,15 @@ def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], ) -> None: """Key what this step generated, from the stop check's host copy.""" + keys = self._runner.prefix_chain_keys(batch.node_name) + if not keys: + return walk = batch.step_context.graph_walk for rid in batch.request_ids: per_rid = outputs.get(rid) if isinstance(per_rid, dict): self._runner.extend_prefix_chains( - rid, batch.node_name, walk, per_rid, + rid, batch.node_name, walk, per_rid, keys=keys, ) def _skip_cached_prefix( From bdf055e9fab9996ef95b00d8c1bf2399a6a46147 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 017/219] rust frames: RESULT_TOKENS, one frame per step for inline client values --- rust/src/graph/frames.rs | 41 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/rust/src/graph/frames.rs b/rust/src/graph/frames.rs index a75234cc6d..6fe4f3dcac 100644 --- a/rust/src/graph/frames.rs +++ b/rust/src/graph/frames.rs @@ -436,6 +436,47 @@ impl ResultTensors<'_> { } } +/// One step's inline emits to the api server: every request's value in one +/// frame, where ResultTensors was one frame (and one tensor) per request. +pub struct ResultTokens<'a> { + pub request_ids: &'a [String], + pub values: &'a [i64], + pub loop_indices: &'a [Option<(Vec, Vec<(Sym, u32)>, u32)>], + pub signal: &'a str, + pub modality: &'a str, +} + +impl ResultTokens<'_> { + pub fn encode(&self, it: &StrToId) -> Vec { + let mut out = Vec::with_capacity(96 + 48 * self.request_ids.len()); + // Untagged like ResultTensors: a declared body type of APIServerMessage. + open_frame(&mut out, "api_msg", "result_tokens", None); + w_map_counted(&mut out, |m| { + w_arr(m.key("request_ids"), self.request_ids.len() as u32); + for r in self.request_ids { + w_str(m.buf(), r); + } + w_arr(m.key("values"), self.values.len() as u32); + for &v in self.values { + w_i64(m.buf(), v); + } + w_arr(m.key("loop_indices"), self.loop_indices.len() as u32); + for li in self.loop_indices { + match li { + Some((order, idxs, fwd)) => { + nested_loop_indices(m.buf(), it, order, idxs, *fwd) + } + None => w_nil(m.buf()), + } + } + w_str(m.key("signal"), self.signal); + w_str(m.key("modality"), self.modality); + w_map(m.key("metadata"), 0); + }); + out + } +} + /// Opens `[outer_tag, {message_type, body}]` and leaves the caller to write the /// body, which is the next value due. `body` is tagged only where its declared /// type is abstract. From 77d32ec40f532a1cb450b8d3b67948ced36dbe19 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 018/219] rust runtime: route an inline client signal without staging, refs or per-request frames --- rust/src/graph/runtime.rs | 106 +++++++++++++++++++++++++++++++++++--- 1 file changed, 98 insertions(+), 8 deletions(-) diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index 305a308780..1af202fc52 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -1076,6 +1076,8 @@ pub struct SendPlan { #[pyo3(get)] pub persist: Vec<(u32, String, Vec)>, /// (rid, signal, modality, tensors) -- sliced, as to_workers is. pub emit: Vec<(u32, String, String, Vec)>, + /// (rid, signal, modality, value): emitted inline, one frame per step. + pub emit_inline: Vec<(u32, String, String, i64)>, /// (rid, finished worker graph ids, is_first_tp_rank) #[pyo3(get)] pub completed: Vec<(u32, Vec, bool)>, /// The pre-completion loop context, snapshotted at route time. @@ -1100,6 +1102,12 @@ pub struct RouteArg { #[pyo3(item)] rids: Vec, #[pyo3(item)] tensors: Vec, #[pyo3(item)] num_tensors: Vec, + /// One output signal whose per-request value is a scalar the worker + /// already holds on the host (a sampled token). Its EMIT_TO_CLIENT edge + /// then rides inline in one frame per step instead of as a tensor per + /// request; `inline_values` are in `rids` order, or empty for none. + #[pyo3(item)] inline_signal: Option, + #[pyo3(item)] inline_values: Vec, } /// `RouteOutput`. @@ -1171,6 +1179,9 @@ pub struct Completion { /// have to ask for it per rid, convert it to strings on the way out, and /// hand it back to be re-interned. pub nested: FxHashMap, Vec<(Sym, u32)>, u32)>, + /// Per rid: (signal, modality, value) of the emit edge that rides inline + /// instead of as a tensor. Empty unless the route carried inline values. + pub emit_inline: FxHashMap, } /// `SpeculationPrepInput`. @@ -2197,6 +2208,16 @@ impl GraphRuntime { let mut nested_snapshot: FxHashMap, Vec<(Sym, u32)>, u32)> = FxHashMap::default(); let mut cursor = 0usize; + // The emit edge of this signal carries its value inline: no staging, + // no per-request frame, no reference for the api server to release. + // A persisted signal keeps the tensor: the conductor reads it. + let inline_sym: Option = input + .inline_signal + .as_deref() + .and_then(|s| self.interner.get(s)); + let inline_on = inline_sym.is_some() + && input.inline_values.len() == input.rids.len(); + let mut emit_inline: FxHashMap = FxHashMap::default(); for (i, &rid) in input.rids.iter().enumerate() { // Slice this rid's uuids back out of the flat, rid-major layout. @@ -2546,9 +2567,19 @@ impl GraphRuntime { // A tensor handled locally is not staged for a remote read; // Python leaves streaming_local and the locally-ingested edge // out of the register set for the same reason. + let inline_edge = inline_on + && matches!(e.dest, Dest::EmitToClient) + && Some(e.name) == inline_sym + && !e.persist; + if inline_edge { + emit_inline.insert( + rid, (e.name, e.modality, input.inline_values[i]), + ); + } let remote = (e.persist && has_remote) || e.declined_local || (!is_local + && !inline_edge && matches!(e.dest, Dest::External(_) | Dest::EmitToClient)); if remote { for t in &e.tensors { @@ -2603,7 +2634,11 @@ impl GraphRuntime { // then nothing counted it and it goes on the wire. Dest::Local(_) if e.declined_local => 1, Dest::Local(_) | Dest::Empty => 0, - Dest::EmitToClient => 1, + // An inline emit holds no reference: nothing reads + // the tensor on the api server's behalf. + Dest::EmitToClient => i64::from( + !(inline_on && Some(e.name) == inline_sym && !e.persist), + ), // One reference per destination WORKER, minus this // one when the edge was already ingested into a local // graph -- that copy is in local_counts. Python does @@ -2713,6 +2748,7 @@ impl GraphRuntime { first_tp_rank, speculative, nested: nested_snapshot, + emit_inline, }, ); Ok(out) @@ -3007,6 +3043,49 @@ impl GraphRuntime { self.dispatch("api_server", &bytes)?; } + // Inline emits: one RESULT_TOKENS frame per (signal, modality) + // per step, with the same per-request bookkeeping as a tensor + // emit so the completion report names the signal and its loop + // context. Usually one group. + if !plan.emit_inline.is_empty() { + let mut groups: Vec<((String, String), Vec<(u32, i64)>)> = Vec::new(); + for (rid, signal, modality, value) in plan.emit_inline { + let key = (signal, modality); + match groups.iter_mut().find(|(k, _)| *k == key) { + Some((_, items)) => items.push((rid, value)), + None => groups.push((key, vec![(rid, value)])), + } + } + for ((signal, modality), items) in groups { + let sig = self.interner.intern(&signal); + let mut request_ids: Vec = Vec::with_capacity(items.len()); + let mut values: Vec = Vec::with_capacity(items.len()); + let mut loop_indices = Vec::with_capacity(items.len()); + for (rid, value) in items { + let request_id = self.rid_name(rid)?; + let idx = nested.get(&rid).cloned(); + if let Some(info) = self.requests[rid as usize].as_mut() { + info.pending.output_signals.push(sig); + if let Some(i) = idx.clone() { + info.pending.set_loop_indices(sig, i); + } + } + request_ids.push(request_id); + values.push(value); + loop_indices.push(idx); + } + let bytes = frames::ResultTokens { + request_ids: &request_ids, + values: &values, + loop_indices: &loop_indices, + signal: &signal, + modality: &modality, + } + .encode(&self.interner); + self.dispatch("api_server", &bytes)?; + } + } + for (rid, signal, uuids) in plan.persist { let sig = self.interner.intern(&signal); if let Some(info) = self.requests[rid as usize].as_mut() { @@ -3065,9 +3144,10 @@ impl GraphRuntime { /// itself are in `frames.rs` (INPUT_SIGNALS, WORKER_GRAPHS_DONE, /// STOP_LOOPS). fn take_send_plan(&mut self, completion_id: u64) -> PyResult { - let c = self.completions.remove(&completion_id).ok_or_else(|| { + let mut c = self.completions.remove(&completion_id).ok_or_else(|| { PyValueError::new_err(format!("unknown completion {completion_id}")) })?; + let emit_inline = std::mem::take(&mut c.emit_inline); let g = self.graphs[c.wg as usize].clone(); let mut plan = SendPlan { partition: c.partition.clone(), @@ -3113,12 +3193,22 @@ impl GraphRuntime { for e in edges { let name = self.interner.name(e.name).to_string(); match e.dest { - Dest::EmitToClient => plan.emit.push(( - rid, - name.clone(), - self.interner.name(e.modality).to_string(), - e.tensors.clone(), - )), + Dest::EmitToClient => match emit_inline.get(&rid) { + Some(&(sig, modality, value)) if sig == e.name => { + plan.emit_inline.push(( + rid, + name.clone(), + self.interner.name(modality).to_string(), + value, + )); + } + _ => plan.emit.push(( + rid, + name.clone(), + self.interner.name(e.modality).to_string(), + e.tensors.clone(), + )), + }, Dest::External(dest) => { // `Sym::MAX` is the fanout's marker for a destination // with no sharding group, i.e. one no worker runs. From 77aab0c43c74b6700197e579ee5169643b93caab Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 019/219] graph runtime: RouteInput.inline_signal/inline_values and supports_inline_emit --- mstar/graph/runtime/base.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mstar/graph/runtime/base.py b/mstar/graph/runtime/base.py index 4b9edf8da7..2f77e613d2 100644 --- a/mstar/graph/runtime/base.py +++ b/mstar/graph/runtime/base.py @@ -427,6 +427,13 @@ class RouteInput(NamedTuple): # rid i's signal s. The runtime looks the metadata up in the store. tensors: list[int] num_tensors: list[int] + # One output signal whose per-rid value is a host scalar the worker + # already has (the sampled token): its EMIT_TO_CLIENT edge then rides + # inline in one frame per step instead of as a tensor per request. The + # values are in wg_ids.keys() order; empty means none. A runtime without + # `supports_inline_emit` ignores both. + inline_signal: str | None = None + inline_values: tuple[int, ...] | list[int] = () class RouteOutput(NamedTuple): @@ -492,6 +499,11 @@ class GraphRuntime(ABC): forwards opaquely) and stream buffers (which hold tensors). """ + # Whether `RouteInput.inline_signal` is honoured: an emit edge sent as one + # RESULT_TOKENS frame per step. A runtime that does not keeps the + # per-request tensor path, which stays correct. + supports_inline_emit: bool = False + # --------- Bookkeeping ---------- @abstractmethod def set_node_metadata( From db4def888a4934de471c4f989fe6d69edc87ff4e Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 020/219] rust runtime wrapper: pass inline values, report supports_inline_emit --- mstar/graph/runtime/rust.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/graph/runtime/rust.py b/mstar/graph/runtime/rust.py index 821a792ebe..386d7dbc9c 100644 --- a/mstar/graph/runtime/rust.py +++ b/mstar/graph/runtime/rust.py @@ -266,6 +266,8 @@ def __init__( self._all_wg_ids_to_nodes = all_wg_ids_to_nodes self._sharding: dict[int, ShardingConfig] = {} + supports_inline_emit = True + # --------- Bookkeeping ---------- def set_node_metadata( @@ -551,6 +553,8 @@ def complete_and_route_batch( "rids": list(input.wg_ids.keys), "tensors": list(input.tensors), "num_tensors": list(input.num_tensors), + "inline_signal": input.inline_signal, + "inline_values": list(input.inline_values), }) return RouteOutput( completion_id=out.completion_id, From 95006f75b3bbdbd85dfe5f1325e2fe5efbfc15b2 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 021/219] api_server: ResultTokens message type --- mstar/api_server/request_types.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/mstar/api_server/request_types.py b/mstar/api_server/request_types.py index 23894cdd8c..ba071ab3e9 100644 --- a/mstar/api_server/request_types.py +++ b/mstar/api_server/request_types.py @@ -25,6 +25,20 @@ class ResultTensors: metadata: dict = field(default_factory=dict) +@dataclass +class ResultTokens: + """One step's client-facing scalars (sampled tokens), every request in one + frame. Replaces one ResultTensors per request for an output whose value + the worker already holds on the host: no transport tensor, no ack.""" + request_ids: list[str] + values: list[int] + # per request, the loop context of the step that produced the value + loop_indices: list[NestedLoopIndices | None] + signal: str + modality: str + metadata: dict = field(default_factory=dict) + + @dataclass class RequestComplete: """Signals that a request has finished processing.""" @@ -55,9 +69,9 @@ class RequestFailed: @dataclass class APIServerMessage: """Envelope for messages received by the API server.""" - # "result_tensors" | "request_complete" | "request_failed" | "setup_done" + # "result_tensors" | "result_tokens" | "request_complete" | "request_failed" | "setup_done" message_type: str - body: ResultTensors | RequestComplete | RequestFailed | None = None # None for setup_done + body: ResultTensors | ResultTokens | RequestComplete | RequestFailed | None = None # None for setup_done @dataclass From 6661a71705ac3d9250c99f635f5489d3f30a6ee8 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 022/219] wire: register result_tokens --- mstar/communication/wire_types.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mstar/communication/wire_types.py b/mstar/communication/wire_types.py index 5b94f6fef9..08093ba951 100644 --- a/mstar/communication/wire_types.py +++ b/mstar/communication/wire_types.py @@ -16,6 +16,7 @@ RequestFailed, ResultChunk, ResultTensors, + ResultTokens, ) from mstar.communication.wire import _set_polymorphic, register from mstar.conductor.request_info import CurrentForwardPassInfo @@ -74,6 +75,7 @@ "fail_requests": FailRequests, # api-server bodies "result_tensors": ResultTensors, + "result_tokens": ResultTokens, "request_complete": RequestComplete, "request_failed": RequestFailed, "result_chunk": ResultChunk, From f6e5939f89f1556cb8d4825c70c7748446173b85 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 023/219] api_server: dispatch result_tokens frames to the data worker --- mstar/api_server/entrypoint.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index e388a82bf8..b00e6a51ba 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -531,6 +531,24 @@ def _process_messages(self) -> None: logger.warning("Unexpected message type: %s", type(message)) continue + if message.message_type == "result_tokens": + # One frame for the step: every request's value. + # A request that is gone needs nothing acked back. + body = message.body + with self.request_lock: + for i, rid in enumerate(body.request_ids): + if rid in self.pending_requests: + self.preprocess_worker.new_result_token( + rid, body.values[i], + body.loop_indices[i], + body.signal, body.modality, + ) + elif rid not in self.recently_completed: + logger.warning( + "Token for unknown request %s dropped", rid, + ) + continue + rid = message.body.request_id with self.request_lock: From 27322e4419043e647a992ab7a62137a1fb148327 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:51 -0700 Subject: [PATCH 024/219] data worker: detokenize inline tokens in arrival order, no read and no ack --- mstar/api_server/data_worker.py | 64 +++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 4671d88c87..f34c75a5b2 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -24,6 +24,7 @@ ResultChunk, ResultTensors, ) + from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig @@ -87,6 +88,15 @@ def _video_frame_metadata( "frame_index": frame_index, "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 + modality: str + def _preprocess_loop(**kwargs): @@ -221,6 +231,24 @@ def new_result_tensors(self, input: ResultTensors): ) self.result_tensor_input_queue.put(input) + def new_result_token( + self, request_id: str, value: int, + loop_indices: NestedLoopIndices | None, signal: str, modality: str, + ): + """A value that rode inline: same accounting as a result tensor, no + transport read and nothing to ack.""" + if request_id not in self.output_loop_idxs: + logger.debug("Late token for cleaned-up request %s dropped", request_id) + return + if loop_indices is not None: + self.output_loop_idxs[request_id][signal] = loop_indices.max( + self.output_loop_idxs[request_id].get(signal, None) + ) + self.per_request_reading_tensors[request_id] += 1 + self.result_tensor_input_queue.put( + InlineToken(request_id, value, loop_indices, modality) + ) + def discard_result_tensors(self, input: ResultTensors): """Ack and drop result tensors for an already-removed request. @@ -663,6 +691,36 @@ 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 + 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, + ) + def _queue_completed_output( self, request_id: str, @@ -888,6 +946,12 @@ def run(self): while not self.result_tensor_queue.empty(): did_work = True result = self.result_tensor_queue.get() + if isinstance(result, InlineToken): + # 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) + continue # Draining for teardown: don't start new reads — ack the # tensors back so the producing worker can free its buffers. if result.request_id in self._draining_rids: From 810cf24860044581380376b917b646ae19d64c4d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 025/219] submodule_base: inline_client_signals hook --- mstar/model/submodule_base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index a010dae8b4..73af2ccaf1 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -677,6 +677,15 @@ def filter_batched_output( a fixed key set for graph compat, so the filtering happens here.""" return outputs + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: + """Output signals whose client-facing value is one scalar per row that + the stop check already copies to the host, as signal name -> the name + of the row-addressed stop buffer holding it. The worker then sends + those values inline, one frame per step, instead of one tensor per + request. Only for a walk whose emit edge is not persisted.""" + del graph_walk + return {} + def unpack_packed_outputs( self, static_output: dict, From 1761cde40ee5bc85b12a35b0701a75afd968f2a8 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 026/219] qwen3_5: the decode loop's client token rides inline --- mstar/model/qwen3_5/submodules.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 42b6fe2649..9601d374dc 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -486,6 +486,11 @@ def check_stop( ) return {"decode_loop"} if hit_eos or out_of_budget else set() + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: + # the decode loop's client edge carries the sampled token the stop + # check already has on the host; prefill's new_token edge is persisted + return {"text_inputs": "new_token"} if graph_walk == "decode" else {} + def check_stop_batched( self, request_ids: list[str], From 02a6af980a5296e2607aea06a2317fef688ffacc Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 027/219] engine: inline_client_signals passthrough --- mstar/engine/engine.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index c35d49cc9d..359183f9ef 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -751,6 +751,15 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: batch=batch, model_inputs=node_inputs ) + def inline_client_signals( + self, node_name: str, graph_walk: str, + ) -> dict[str, str]: + """The node's client signals that can ride inline this walk.""" + mgmt = self._submodules.get(node_name) + if mgmt is None: + return {} + return mgmt.submodule.inline_client_signals(graph_walk) + def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], ) -> None: From d57681cf15e84da768625fd67a7441727111734b Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 028/219] worker: hand the stop check's host tokens to the route as inline values --- mstar/worker/worker.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 3ce5b3f973..1aa6de4a0a 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2898,6 +2898,9 @@ def _pp_stage(name: str) -> None: # The graph runtime's own share of postprocess: the routing call. _t_route = _time.perf_counter() if self._phase_period else 0.0 + inline_signal, inline_values = self._inline_emit_values( + engine, batch_N, host_rows, rids, signals, + ) route_output = self._graph_runtime.complete_and_route_batch( RouteInput( partition=batch_N.partition, @@ -2910,6 +2913,8 @@ def _pp_stage(name: str) -> None: ), tensors=flat_uuids, num_tensors=num_tensors, + inline_signal=inline_signal, + inline_values=inline_values, ), ) if self._phase_period: @@ -3048,6 +3053,37 @@ def _pp_stage(name: str) -> None: if self.enable_nvtx: range_pop(synchronize=False) + def _inline_emit_values( + self, engine, batch_N: PendingBatch, host_rows: "HostRows | None", + rids: list[int], signals, + ) -> tuple[str | None, list[int]]: + """The values of a client signal that can ride inline this step. + + Taken off the stop check's host copy, one ``tolist`` for the batch. + Every rid must have a row, or the step keeps the tensor path for the + whole batch: the runtime applies the values positionally. + """ + if host_rows is None or not self._graph_runtime.supports_inline_emit: + return None, [] + candidates = engine.inline_client_signals( + batch_N.node_name, batch_N.graph_walk, + ) + if not candidates: + return None, [] + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + for signal, buffer_name in candidates.items(): + if signal not in signals: + continue + buf = host_rows.buffers.get(buffer_name) + if not torch.is_tensor(buf) or buf.dim() == 0 or buf.numel() == 0: + continue + values = buf.reshape(buf.shape[0], -1)[:, 0].tolist() + try: + return signal, [values[row_of[rid]] for rid in rids] + except (KeyError, IndexError): + return None, [] + return None, [] + def _get_pinned_d2h_buffer( self, purpose: str, From 21fb6961ce80c7fc553a36f0663a52d370e81f17 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 029/219] test: inline token path (wire, data worker, worker) --- test/modular/test_inline_tokens.py | 192 +++++++++++++++++++++++++++++ 1 file changed, 192 insertions(+) create mode 100644 test/modular/test_inline_tokens.py diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py new file mode 100644 index 0000000000..a20cb35938 --- /dev/null +++ b/test/modular/test_inline_tokens.py @@ -0,0 +1,192 @@ +"""The inline token path: a decode step's sampled tokens ride to the api server +as one RESULT_TOKENS frame (values, not tensors). Covers the wire type, the +api server's data worker, and the worker's value extraction.""" +from __future__ import annotations + +import queue +import threading + +import torch + +from mstar.api_server.data_worker import InlineToken, PreprocessWorker, PreprocessWorkerThread +from mstar.api_server.request_types import APIServerMessage, ResultChunk, ResultTokens +from mstar.communication import wire +from mstar.graph.loop_indices import NestedLoopIndices +from mstar.model.submodule_base import HostRows +from mstar.worker.worker import Worker + +# -- wire -------------------------------------------------------------------- + + +def _loop(i): + return NestedLoopIndices(loop_name_order=["decode_loop"], loop_indices={"decode_loop": i}, wg_fwd_pass_idx=i + 1) + + +def test_result_tokens_round_trips_through_the_codec(): + msg = APIServerMessage( + message_type="result_tokens", + body=ResultTokens( + request_ids=["r1", "r2", "r3"], values=[5, 151645, 7], + loop_indices=[_loop(3), None, _loop(9)], + signal="text_inputs", modality="text", + ), + ) + back = wire.decode(wire.encode(msg)) + assert back.message_type == "result_tokens" + assert isinstance(back.body, ResultTokens) + assert back.body.request_ids == ["r1", "r2", "r3"] + assert back.body.values == [5, 151645, 7] + assert back.body.loop_indices[1] is None + assert back.body.loop_indices[0].loop_indices == {"decode_loop": 3} + assert back.body.loop_indices[2].wg_fwd_pass_idx == 10 + assert back.body.signal == "text_inputs" and back.body.modality == "text" + + +# -- api server -------------------------------------------------------------- + + +class _Model: + def postprocess(self, output, modality, request_kwargs=None): + assert modality == "text" + return f"<{int(output[0])}>".encode() + + +class _Comm: + def get_all_new_messages(self): + return [] + + def send(self, entity, msg): + pass + + +class _TM: + def start_read_tensors(self, request_id, graph_edges, graph_walk=None): + return [] + + def get_ready_tensors(self): + return {} + + def has_inflight_reads(self, request_id): + return False + + def cleanup_request(self, request_id): + pass + + def force_cleanup_request(self, request_id): + pass + + def ack_unread_tensors(self, request_id, graph_edges): + pass + + +def _thread(): + stop = threading.Event() + worker = PreprocessWorkerThread( + in_queue=queue.Queue(), result_tensor_queue=queue.Queue(), out_queue=queue.Queue(), + profile_queue=queue.Queue(), cleanup_request_queue=queue.Queue(), + abort_request_queue=queue.Queue(), reads_done_queue=queue.Queue(), + discard_tensor_queue=queue.Queue(), stop_event=stop, communicator=_Comm(), + tensor_manager=_TM(), model=_Model(), + ) + return worker, stop + + +def test_inline_tokens_are_detokenized_in_order_without_a_read(): + worker, stop = _thread() + for i, v in enumerate((11, 12, 13)): + worker.result_tensor_queue.put(InlineToken("req", v, _loop(i), "text")) + worker.result_tensor_queue.put(InlineToken("other", 99, None, "text")) + thread = threading.Thread(target=worker.run) + thread.start() + try: + chunks = [worker.out_queue.get(timeout=5) for _ in range(4)] + finally: + stop.set() + thread.join(timeout=5) + assert all(isinstance(c, ResultChunk) for c in chunks) + assert [c.data for c in chunks if c.request_id == "req"] == [b"<11>", b"<12>", b"<13>"] + assert [c.data for c in chunks if c.request_id == "other"] == [b"<99>"] + state = worker.request_output_state["req"] + assert state.next_sequence == 3 and state.next_emit == 3 and not state.pending + + +def test_a_draining_request_drops_its_inline_tokens(): + worker, stop = _thread() + worker._draining_rids.add("gone") + worker.result_tensor_queue.put(InlineToken("gone", 1, None, "text")) + worker.result_tensor_queue.put(InlineToken("live", 2, None, "text")) + thread = threading.Thread(target=worker.run) + thread.start() + try: + chunk = worker.out_queue.get(timeout=5) + finally: + stop.set() + thread.join(timeout=5) + assert chunk.request_id == "live" and chunk.data == b"<2>" + assert worker.out_queue.empty() + + +def test_the_facade_accounts_an_inline_token_like_a_tensor(): + facade = PreprocessWorker.__new__(PreprocessWorker) + facade.result_tensor_input_queue = queue.Queue() + facade.per_request_reading_tensors = {"r1": 0} + facade.output_loop_idxs = {"r1": {}} + facade.new_result_token("r1", 42, _loop(4), "text_inputs", "text") + facade.new_result_token("r1", 43, _loop(5), "text_inputs", "text") + facade.new_result_token("unknown", 1, _loop(0), "text_inputs", "text") # dropped, no raise + assert facade.per_request_reading_tensors["r1"] == 2 + assert facade.output_loop_idxs["r1"]["text_inputs"].loop_indices == {"decode_loop": 5} + items = [facade.result_tensor_input_queue.get_nowait() for _ in range(2)] + assert [i.value for i in items] == [42, 43] + assert facade.result_tensor_input_queue.empty() + + +# -- worker ------------------------------------------------------------------ + + +class _Runtime: + supports_inline_emit = True + + +class _Engine: + def __init__(self, signals): + self._signals = signals + + def inline_client_signals(self, node_name, graph_walk): + return self._signals + + +class _Batch: + node_name = "LLM" + graph_walk = "decode" + + +def _worker(runtime=_Runtime()): + w = Worker.__new__(Worker) + w._graph_runtime = runtime + return w + + +def test_inline_values_follow_the_batch_order_off_the_host_rows(): + rows = HostRows((7, 8, 9), {"new_token": torch.tensor([[70], [80], [90]])}) + sig, vals = _worker()._inline_emit_values( + _Engine({"text_inputs": "new_token"}), _Batch(), rows, [9, 7, 8], ["text_inputs", "kv"], + ) + assert sig == "text_inputs" and vals == [90, 70, 80] + + +def test_inline_values_fall_back_when_a_rid_has_no_row_or_nothing_applies(): + rows = HostRows((7, 8), {"new_token": torch.tensor([70, 80])}) + w = _worker() + eng = _Engine({"text_inputs": "new_token"}) + none = (None, []) + assert w._inline_emit_values(eng, _Batch(), rows, [7, 8, 9], ["text_inputs"]) == none + assert w._inline_emit_values(eng, _Batch(), rows, [7], ["other"]) == none + assert w._inline_emit_values(_Engine({}), _Batch(), rows, [7], ["text_inputs"]) == none + assert w._inline_emit_values(eng, _Batch(), None, [7], ["text_inputs"]) == none + + class _NoInline: + supports_inline_emit = False + assert _worker(_NoInline())._inline_emit_values( + eng, _Batch(), rows, [7, 8], ["text_inputs"], + ) == none From a31c952c579a90f469dca87c103db037d510347a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:11:52 -0700 Subject: [PATCH 030/219] test: rust runtime inline emit --- test/rust/test_rust_graph_runtime.py | 107 +++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/test/rust/test_rust_graph_runtime.py b/test/rust/test_rust_graph_runtime.py index a40e6358aa..f5b274006d 100644 --- a/test/rust/test_rust_graph_runtime.py +++ b/test/rust/test_rust_graph_runtime.py @@ -904,3 +904,110 @@ def test_a_local_only_batch_sends_nothing_to_peers(runtime): peers = [e for e, _ in rec.sent if e not in ("conductor", "api_server")] assert peers == [] + + +# --- inline emit ------------------------------------------------------------- + +def _client_loop_runtime(): + """A decode node whose token goes back to itself AND to the client.""" + from mstar.graph.special_destinations import EMIT_TO_CLIENT + section = Sequential(sections=[ + GraphNode( + name="prefill", input_names={"prompt"}, + outputs=[GraphEdge(name="token", next_node="dec")], + ), + Loop( + name="loop", + section=GraphNode( + name="dec", input_names={"token"}, + outputs=[ + GraphEdge(name="token", next_node="dec"), + GraphEdge(name="token", next_node=EMIT_TO_CLIENT, output_modality="text"), + ], + ), + outputs=[], + max_iters=50, + ), + ]) + wg = WorkerGraph(section=section, graph_walks={WALK}, ranks=[0], worker_graph_id=WG_ID) + book = RustTensorBookkeeping() + rt = rust_runtime.RustGraphRuntime( + my_worker_id=WORKER, my_worker_graphs=[wg], + all_wg_ids_to_graph_walks={WG_ID: {WALK}}, + all_wg_ids_to_dyn_loops={WG_ID: {"loop"}}, + all_wg_ids_to_nodes={WG_ID: {"prefill", "dec"}}, + node_to_partition={"prefill": "default", "dec": "default"}, + sharding_config=ShardingConfig(groups=[], tp_enabled_nodes=set(), shard_dim={}), + bookkeeping=book, + ) + rt._bookkeeping_for_test = book + return rt + + +def _run_dec_step(rt, rid, uuid, inline_value=None): + bk = rt._bookkeeping_for_test + bk.put_tensor(uuid, _tensor_info(uuid)) + bk.increment_ref(uuid, 1) # the worker's safety hold + kwargs = {} + if inline_value is not None: + kwargs = dict(inline_signal="token", inline_values=[inline_value]) + return rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], + wg_ids=ParallelList([rid], [WG_ID]), + tensors=[uuid], num_tensors=[1], **kwargs, + )) + + +def _reach_dec(rt, rid, uuid): + rt.ingest_inputs_batch(_ingest_block([rid], [_spec("prompt", "prefill")])) + rt.pop_rids("prefill", WALK, [rid]) + bk = rt._bookkeeping_for_test + bk.put_tensor(uuid, _tensor_info(uuid)) + bk.increment_ref(uuid, 1) + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="prefill", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[uuid], num_tensors=[1], + )) + rt.send_outputs(_send_input(rid, out.completion_id)) + rt.pop_rids("dec", WALK, [rid]) + + +def test_the_rust_runtime_offers_inline_emit(): + assert rust_runtime.RustGraphRuntime.supports_inline_emit is True + assert rust_runtime_base.GraphRuntime.supports_inline_emit is False + + +def test_an_inline_emit_is_not_staged_for_the_client(): + rt = _client_loop_runtime() + rid = _admit(rt) + rt._communicator = _Recorder() + _reach_dec(rt, rid, 100) + # the tensor path stages the client's copy + out = _run_dec_step(rt, rid, 101) + assert 101 in out.register_uuids + rt.send_outputs(_send_input(rid, out.completion_id)) + rt.pop_rids("dec", WALK, [rid]) + # the inline path does not: the value rides in the frame + out = _run_dec_step(rt, rid, 102, inline_value=151645) + assert 102 not in out.register_uuids + rt.send_outputs(_send_input(rid, out.completion_id)) + # the loop-back still fed the node, so the next step is ready + assert rt.pop_rids("dec", WALK, [rid]) is not None + + +def test_inline_values_must_cover_every_rid_or_are_ignored(): + rt = _client_loop_runtime() + rid = _admit(rt) + _reach_dec(rt, rid, 200) + bk = rt._bookkeeping_for_test + bk.put_tensor(201, _tensor_info(201)) + bk.increment_ref(201, 1) + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[201], num_tensors=[1], inline_signal="token", inline_values=[], + )) + assert 201 in out.register_uuids, "no values, so the tensor path" + rt.send_outputs(_send_input(rid, out.completion_id)) From b61b1b594c45294e5112a1270620f88174e7df01 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:16:59 -0700 Subject: [PATCH 031/219] qwen3_5: decode capture buckets up to 128 rows --- mstar/model/qwen3_5/submodules.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 9601d374dc..9655d1240f 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -79,7 +79,11 @@ class LLMSubmodule(ARNodeSubmodule): # Capture rows and a replay's padding rows address the pool's sink and # hold no slot, so these buckets do not size `gdn_state.max_slots`; that is # set by the concurrency a deployment wants (see configs/qwen3_5_*.yaml). - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # The ladder runs to 128 because the largest bucket is also the scheduler's + # cap on rows per decode step: with 32 as the top, 64 requests in flight + # ran as two alternating batches of 32 and each paid the whole per-step + # host cost. A step above 32 rows pads to the next bucket. + DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 48, 64, 96, 128] # Text plus image tokens of a whole prompt. Capped at 4096, not the # processor's 16384, because static buffers are sized by the largest From 51edd0cf05d695125135a12214b33234be3d2a71 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:18:05 -0700 Subject: [PATCH 032/219] api_server: wake the streaming handler on events instead of polling per request --- mstar/api_server/entrypoint.py | 38 +++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index b00e6a51ba..3b347bc40b 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -184,6 +184,21 @@ class PendingRequest: consumed_chunks: int = 0 error: Any | None = None error_status: int = 500 + # Set by the streaming handler on its event loop; the message thread wakes + # it through these when a chunk or the completion lands, so the handler + # sleeps between chunks instead of polling every millisecond per request. + wake: Any | None = None + loop: Any | None = None + + def notify(self) -> None: + """Wake the streaming handler, from any thread.""" + loop, wake = self.loop, self.wake + if loop is None or wake is None: + return + try: + loop.call_soon_threadsafe(wake.set) + except RuntimeError: + pass # the loop is closed: the client is gone def _chunk_to_ndjson_payload(chunk: ResultChunk) -> str: @@ -361,6 +376,7 @@ def _fail_pending_for_dead_conductor(self, exited: str) -> None: req.error = message req.error_status = 503 req.event.set() + req.notify() on_fatal = self.on_fatal if on_fatal is not None: on_fatal() @@ -491,6 +507,7 @@ def _prune_recently_completed(self) -> None: ) req.error_status = 500 req.event.set() + req.notify() # Snapshot the data worker's tx/rx now: the request is done (all # final chunks received), so the worker thread is no longer mutating # this rid's transport state, and we must read it before the hard @@ -577,6 +594,7 @@ def _process_messages(self) -> None: # result is coming, so the alternative is the # blanket request timeout. req.event.set() + req.notify() # The conductor has already dropped this rid, # so no abort is needed — but the data worker # still holds its transport state. Parking the @@ -656,6 +674,7 @@ def _process_messages(self) -> None: req.profile.timing.first_chunk_time = now req.profile.timing.last_chunk_time = now req.chunks.append(result_chunk) + req.notify() if result_chunk.modality == "error": # The data worker failed this request (preprocess, @@ -668,6 +687,7 @@ def _process_messages(self) -> None: (result_chunk.metadata or {}).get("status", 500) ) req.event.set() + req.notify() # Park the rid so _prune_recently_completed reclaims # the data worker's per-request state once the # client lets go of the request. @@ -694,11 +714,22 @@ async def iter_result_chunks(self, request_id: str): """ start = time.time() finished = False + # Woken by the message thread (PendingRequest.notify) when a chunk or + # the completion lands; the timeout below is only a backstop. + wake = asyncio.Event() + with self.request_lock: + req = self.pending_requests.get(request_id) + if req is not None: + req.loop = asyncio.get_running_loop() + req.wake = wake try: while True: if time.time() - start > self.timeout_seconds: raise HTTPException(status_code=500, detail="Request timed out") + # Cleared before the read: a notification that lands after + # this point stays set and the wait below returns at once. + wake.clear() new_chunks: list[ResultChunk] = [] done = False with self.request_lock: @@ -747,7 +778,12 @@ async def iter_result_chunks(self, request_id: str): finished = True break - await asyncio.sleep(0.001) + if new_chunks: + continue # more may already be waiting + try: + await asyncio.wait_for(wake.wait(), timeout=0.05) + except asyncio.TimeoutError: + pass finally: if not finished: self.abort_request(request_id) From 8d21448949c3622e28cae12ecb354e5f0a346a96 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:18:05 -0700 Subject: [PATCH 033/219] test: a chunk from the message thread wakes the stream --- test/modular/test_api_result_delivery.py | 39 ++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/test/modular/test_api_result_delivery.py b/test/modular/test_api_result_delivery.py index 7057e3908a..92901871f5 100644 --- a/test/modular/test_api_result_delivery.py +++ b/test/modular/test_api_result_delivery.py @@ -235,3 +235,42 @@ async def _collect(): chunks = asyncio.run(_collect()) assert [c.modality for c in chunks] == ["text", "error"] assert chunks[-1].metadata["status"] == 500 + + +def test_a_chunk_from_another_thread_wakes_the_stream_at_once(): + """The handler sleeps on an event the message thread sets, so a chunk is + delivered without waiting out a poll interval (the backstop is 50 ms).""" + pw = _StubPreprocessWorker() + server = _api_server_stub(pw) + req = _pending_request() + server.pending_requests["r4"] = req + + async def _collect(): + loop = asyncio.get_running_loop() + times = [] + + def producer(): + time.sleep(0.2) + for i in range(3): + time.sleep(0.15) + with server.request_lock: + req.chunks.append(ResultChunk(request_id="r4", modality="text", data=f"c{i}".encode())) + req.notify() + times.append(time.perf_counter()) + time.sleep(0.15) + with server.request_lock: + req.event.set() + req.notify() + + threading.Thread(target=producer, daemon=True).start() + got = [] + async for chunk in server.iter_result_chunks("r4"): + got.append((chunk.data, time.perf_counter())) + del loop + return got, times + + got, times = asyncio.run(_collect()) + assert [d for d, _ in got] == [b"c0", b"c1", b"c2"] + # each chunk reached the handler within the 50 ms backstop, not a poll later + for (_, t_got), t_put in zip(got, times, strict=True): + assert t_got - t_put < 0.05, f"chunk took {1000 * (t_got - t_put):.1f} ms" From 85c53671672cb0171aecced6bb69d30169648e73 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:19:43 -0700 Subject: [PATCH 034/219] submodule_base: step_is_input_free hook --- mstar/model/submodule_base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 73af2ccaf1..de9fc00f2f 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -677,6 +677,15 @@ def filter_batched_output( a fixed key set for graph compat, so the filtering happens here.""" return outputs + def step_is_input_free(self, graph_walk: str) -> bool: + """Whether ``declare_step`` for this walk reads nothing from the inputs + but their spans, and every real row has the span of the template row. + Then the step a pre-plan declared over the bucket's template rows is + the step the forward would declare, and the engine runs it instead of + declaring and admitting a second time.""" + del graph_walk + return False + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: """Output signals whose client-facing value is one scalar per row that the stop check already copies to the host, as signal name -> the name From ba56be0db8a9ea04b94dd6f4fde52c4d18b74a5e Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:19:43 -0700 Subject: [PATCH 035/219] qwen3_5: decode steps declare from spans alone --- mstar/model/qwen3_5/submodules.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 9655d1240f..a29ce50e21 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -490,6 +490,11 @@ def check_stop( ) return {"decode_loop"} if hit_eos or out_of_budget else set() + def step_is_input_free(self, graph_walk: str) -> bool: + # decode rows are one token each, like the template; the prefill walks + # feed `prefill_tokens` from the inputs + return graph_walk == "decode" + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: # the decode loop's client edge carries the sampled token the stop # check already has on the host; prefill's new_token edge is persisted From 2ce4dacf2b98024ac320e4c174ee13872488eb15 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:19:43 -0700 Subject: [PATCH 036/219] engine: run the pre-plan's declared step for input-free walks --- mstar/engine/engine.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 359183f9ef..5a9d63f45d 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1798,7 +1798,14 @@ def _pre_plan_for_batch( ) if step is None: return False - if batch.inputs is not None: + if batch.inputs is not None or ( + not submodule_mgmt.piecewise_runners + and submodule_mgmt.submodule.step_is_input_free( + batch.step_context.graph_walk, + ) + ): + # exec drives this step as declared here (its admit ran as + # pre_admit), instead of declaring and admitting again batch.step = step batch.step_context.is_preplan = True @@ -1841,6 +1848,10 @@ def reset_pre_plan_for_batch(self, batch: ExecutingBatch | None = None) -> None: mark(f"engine.preplan_discarded.{batch.node_name}") batch.preplan_event = None batch.preplanned_rids = None + if batch.inputs is None: + # declared over template rows for a step that is now stale: + # exec declares again over the real ones + batch.step = None lease = batch.step_context.slot_lease cg_runner = self._submodules[batch.node_name].cuda_graph_runner if lease is not None and cg_runner is not None: From 6557266cef847bb58e81c288638aeacf7cc78272 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:20:52 -0700 Subject: [PATCH 037/219] docs: MSTAR_LAZY_PUBLISH --- docs/environment_variables.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 0c09c1aec3..878435e580 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -295,6 +295,14 @@ Worker scheduling - Pre-plan the speculative batch's attention on a dedicated thread while the previous replay runs. ``0`` plans inline on the GPU thread. + * - ``MSTAR_LAZY_PUBLISH`` + - ``1`` + - Publish resource state (KV pages, positions) only for the requests + whose frame or completion carries it: ``finalize_batch`` keeps a + cheap per-request snapshot and the worker finishes it on demand. + ``0`` builds the full publication for every request on every step, + the old behaviour; it was a quarter of the GPU thread's host time at + batch 32 in a decode loop that never reads it. * - ``MSTAR_MAX_CONSECUTIVE_SPEC_STEPS`` - ``1024`` - Cap on back-to-back speculative steps before the leader yields to From 080065cbd5033cc62c47503b3979c6dd15008800 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:57:10 -0700 Subject: [PATCH 038/219] worker: build and pre-plan N+2 while N+1 runs, between N's routing and send --- mstar/worker/worker.py | 355 ++++++++++++++++++++++++++++------------- 1 file changed, 245 insertions(+), 110 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 1aa6de4a0a..70f52dd3a7 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -206,6 +206,21 @@ class EvictionPolicy(Enum): # resource to consult. +@dataclass +class _PostprocessState: + """What the first half of post-processing (stops, store, route) hands the + send stage.""" + cpu_outputs: dict + host_rows: "HostRows | None" + buffered_publish: set[int] + route_output: object = None + rids: list = field(default_factory=list) + signals: object = None + flat_rids: list = field(default_factory=list) + flat_uuids: list = field(default_factory=list) + signal_idxs: object = None + + class Worker: """ Real worker that integrates RequestStateManager, EngineManager, @@ -433,6 +448,9 @@ def __init__( # speculation; REMOVE_REQUEST for these is deferred. # _pending_removes: deferred REMOVE_REQUESTs. self._in_flight_rids: set[int] = set() + # rids of a speculation built and armed ahead (early speculation); + # kept in flight so a REMOVE for them is deferred until it ran + self._armed_spec_rids: set[int] = set() self._pending_removes: set[int] = set() # step seq -> rids rank 0 released after that step, held until this rank @@ -2670,6 +2688,23 @@ def _postprocess_batch( self, batch_N: PendingBatch, outputs: BatchedModelOutput, ): + """Both halves back to back: the stop decisions, then the routing.""" + state = self._postprocess_batch_stops(batch_N, outputs) + if state is not None: + self._postprocess_batch_send(batch_N, outputs, state) + + def _postprocess_batch_stops( + self, batch_N: PendingBatch, + outputs: BatchedModelOutput, + ) -> "_PostprocessState | None": + """Post-processing N up to and including the routing: reclaim its + inputs, apply the stops decided a step ago, pull the sampled tokens to + the host, decide this step's stops, publish what stopped, store and + route the outputs (which advances the loop counters). Returns None when + nothing is left to send. The send stage is ``_postprocess_batch_send``; + the split lets the loop build the next speculation in between, once + N's stops and loop counters are settled and while N+1 runs. + """ # Stage stopwatch named like the NVTX ranges below; a clock stamp, not # a nested `_span`, because of the two early returns. _pp_t = _time.perf_counter() if self._phase_period else 0.0 @@ -2717,7 +2752,7 @@ def _pp_stage(name: str) -> None: if not valid_rids: if self.enable_nvtx: range_pop(synchronize=False) - return + return None # pending stops are only needed for one iteration, so can be cleared now self._graph_runtime.clear_pending_loop_stops() @@ -2808,7 +2843,7 @@ def _pp_stage(name: str) -> None: if not batch_N.node_batch.request_ids: if self.enable_nvtx: range_pop(synchronize=False) - return + return None _pp_stage("check_stop") if self.enable_nvtx: @@ -2855,6 +2890,7 @@ def _pp_stage(name: str) -> None: _pp_stage("stop_loops") if self.enable_nvtx: range_pop(synchronize=False) + if self.enable_nvtx: range_push("worker.postprocess.route_outputs", synchronize=False) # Store this batch's output tensors, then hand the runtime their # uuids. The store keeps the descriptors, so routing needs no @@ -2941,6 +2977,36 @@ def _pp_stage(name: str) -> None: self._register_outputs(route_output) _pp_stage("register_outputs") + return _PostprocessState( + cpu_outputs=cpu_outputs, host_rows=host_rows, + buffered_publish=buffered_publish, route_output=route_output, + rids=rids, signals=signals, flat_rids=flat_rids, + flat_uuids=flat_uuids, signal_idxs=signal_idxs, + ) + + def _postprocess_batch_send( + self, batch_N: PendingBatch, + outputs: BatchedModelOutput, + state: "_PostprocessState", + ) -> None: + """The last stage of post-processing N: the frames to peers, the + api server and the conductor. Split off so the loop can build the next + speculation before it, once N's routing advanced the loop counters.""" + del outputs + route_output = state.route_output + rids, signals = state.rids, state.signals + flat_rids, flat_uuids, signal_idxs = state.flat_rids, state.flat_uuids, state.signal_idxs + buffered_publish = state.buffered_publish + engine = self.engine_manager.get_engine(batch_N.node_name) + _pp_t = _time.perf_counter() if self._phase_period else 0.0 + def _pp_stage(name: str) -> None: + nonlocal _pp_t + if not self._phase_period: + return + now = _time.perf_counter() + self._phase_buf[f"worker.postprocess.{name}"].append(now - _pp_t) + _pp_t = now + # send outputs if self.enable_nvtx: range_pop(synchronize=False) @@ -3307,6 +3373,7 @@ def _handle_main_loop_error( in_flight: "tuple[PendingBatch | None, ...]", batch: ScheduledBatch | None, speculation: "Speculation | None" = None, + next_speculation: "Speculation | None" = None, ) -> None: """Fail everything the crashed iteration touched and drain its futures. @@ -3350,7 +3417,19 @@ def _handle_main_loop_error( if batch is not None: failed_rids.update(batch.request_to_worker_graph) self._clear_in_flight_flag(batch) - if speculation is not None and speculation.plan_future is not None: + self._armed_spec_rids = set() + for spec in (speculation, next_speculation): + if spec is None: + continue + self._unwind_speculation(spec, failed_rids) + + self._fail_requests({rid: f"Error in worker: {err}" for rid in failed_rids}) + + def _unwind_speculation( + self, speculation: "Speculation", failed_rids: set[int], + ) -> None: + """Drain an armed speculation's pre-plan and give its fresh rids back.""" + if speculation.plan_future is not None: # Armed but never submitted: the plan thread may still be staging # a pre-plan for a step that will now never run. Drain it, then # drop the stage, or the next step to lease that slot would find @@ -3383,8 +3462,6 @@ def _handle_main_loop_error( sb.node_name, list(fresh), list(fresh.values()), ) - self._fail_requests({rid: f"Error in worker: {err}" for rid in failed_rids}) - def _fail_requests(self, errors: dict[int, str]) -> None: """Report requests this worker can no longer serve to the conductor. @@ -3540,11 +3617,120 @@ def run(self) -> None: ) consecutive_spec_steps = 0 yield_away_from_target: tuple[str, str] | None = None + # Build N+2 while N+1 runs, between the two halves of N's + # post-processing: its pre-plan then overlaps the routing half rather + # than the next launch. TP1 nodes only; a parallel node keeps the head + # broadcast in the order the followers settle it. + early_spec = os.environ.get("MSTAR_EARLY_SPEC", "1") == "1" + next_speculation: Speculation | None = None + next_yield_away: tuple[str, str] | None = None + from mstar.utils.profiler import range_pop, range_push + + def _build_speculation( + pending: PendingBatch | None, consecutive: int, + ) -> tuple[Speculation | None, tuple[str, str] | None]: + """Speculate the step after ``pending`` (N+1 from N), or a + yield-away batch when fairness says so. Returns the speculation + and the (node, walk) to exclude on the non-speculative path.""" + speculation = None + yield_away_from_target = None + batch = None + admit_refused = ( + pending is not None + and pending.node_batch.admit_error is not None + ) + + if pending is not None and self._can_speculate(pending.batch): + # Fairness check (peek-based, replaces the old iter- + # counter cap): only break the spec chain when there's + # another (node, walk) actually ready to schedule on + # this worker. On single-walk workers (Orpheus LLM, + # Orpheus SNAC) this returns False and we always speculate. + must_yield_for_fairness = ( + spec_peek_for_fairness + and consecutive >= 1 + and self.scheduler.has_ready_excluding( + self.request_state, + (pending.node_name, pending.graph_walk), + ) + ) + must_yield_away = ( + consecutive >= max_consecutive_spec + or must_yield_for_fairness + ) + if not must_yield_away and not admit_refused: + if self.enable_nvtx: + range_push("worker.speculate", synchronize=False) + _t0 = _time.perf_counter() if phase_period else 0.0 + speculation = self._try_speculate_next(pending) + if phase_period: + _phase_record("speculate", _time.perf_counter() - _t0) + if self.enable_nvtx: + range_pop(synchronize=False) + if speculation is not None: + # Broadcast the head now, during forward N, so + # followers build it during theirs (-1: not parallel). + speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( + speculation.node_batch, + speculative=True, spec_from_seq=pending.tp_seq, + ) + if self._tp_lead_needs_marker(pending, speculation): + self._broadcast_tp_nospec(pending) + # ``yield_away_from_target`` stays None when the admit + # refused, so the non-speculative path below schedules + # without the fairness exclusion — it should be free to pick + # up the batch ``_handle_admit_failure`` just pushed back. + if speculation is None and not admit_refused: + yield_away_from_target = ( + pending.node_name, + pending.graph_walk, + ) if must_yield_away else None + with self._span("worker.schedule_yield_away"): + batch = self.scheduler.get_next_batch( + self.request_state, + exclude_target=yield_away_from_target, + ) + if batch is not None: + node_batch = self._build_executing_batch(batch) + batch_partition = self.request_state.get_partition_for_node(batch.node_name) + logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") + speculation = Speculation( + scheduled_batch=batch, + node_batch=node_batch, + consumed_edges=set(), + continuing_rids=set(), # n/a + partition=batch_partition, + is_new_iter=False, + is_same_node=False, + is_yield_away=True + ) + + # A leader stamps the seq it sends; a follower's + # batch keeps the one it came off the FIFO with. + ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) + speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq + + return speculation, yield_away_from_target + + def _arm_speculation( + spec: Speculation, pending_for: PendingBatch | None = None, + ) -> None: + # Hand N+1 to the plan thread NOW. It waits on N's + # commit event, then reserves a slot and pre-plans — + # while the main thread sits in await_gpu with the GIL + # released and N's kernels are still running. + if plan_executor is not None: + spec.plan_future = plan_executor.submit( + self._preplan_spec, + pending if pending_for is None else pending_for, spec, + ) def _set_pending(p: PendingBatch): nonlocal pending pending = p - self._in_flight_rids = set(p.batch.request_to_worker_graph) if p else set() + self._in_flight_rids = ( + set(p.batch.request_to_worker_graph) if p else set() + ) | self._armed_spec_rids # Per-phase wall-clock instrumentation, gated by MSTAR_PHASE_TIMING. # When enabled, every Nth speculative iter logs a histogram so we can @@ -3599,7 +3785,6 @@ def _phase_flush() -> None: speculation: Speculation | None = None while True: - from mstar.utils.profiler import range_pop, range_push try: batch = None spec_pending = None @@ -3645,105 +3830,20 @@ def _phase_flush() -> None: # Only when (a) there's a pending step and (b) it's AR-engine. # For non-AR or non-loop-body steps, falls through to the # non-speculative path below (drain, then schedule). - speculation = None - yield_away_from_target = None - - # Build nothing on a refusal; the fence above is what makes it - # knowable here. Ordinary speculation would be dropped anyway, - # but a yield-away batch survives ``_maybe_clear_spec`` and is - # already broadcast — so the eviction below would land after its - # head went out, and the follower would admit it a delta behind. - # Skipping keeps that eviction in the gap between steps. - # - # The marker below stays outside this: a follower settles N on - # exactly one of {head, marker}, and it moves no pages. - admit_refused = ( - pending is not None - and pending.node_batch.admit_error is not None - ) - - if pending is not None and self._can_speculate(pending.batch): - # Fairness check (peek-based, replaces the old iter- - # counter cap): only break the spec chain when there's - # another (node, walk) actually ready to schedule on - # this worker. On single-walk workers (Orpheus LLM, - # Orpheus SNAC) this returns False and we always speculate. - must_yield_for_fairness = ( - spec_peek_for_fairness - and consecutive_spec_steps >= 1 - and self.scheduler.has_ready_excluding( - self.request_state, - (pending.node_name, pending.graph_walk), - ) - ) - must_yield_away = ( - consecutive_spec_steps >= max_consecutive_spec - or must_yield_for_fairness + if next_speculation is not None: + # Built and armed during the previous iteration's + # post-processing (early speculation, below). + speculation = next_speculation + yield_away_from_target = next_yield_away + next_speculation = next_yield_away = None + self._armed_spec_rids = set() + else: + speculation, yield_away_from_target = _build_speculation( + pending, consecutive_spec_steps, ) - if not must_yield_away and not admit_refused: - if self.enable_nvtx: - range_push("worker.speculate", synchronize=False) - _t0 = _time.perf_counter() if phase_period else 0.0 - speculation = self._try_speculate_next(pending) - if phase_period: - _phase_record("speculate", _time.perf_counter() - _t0) - if self.enable_nvtx: - range_pop(synchronize=False) - if speculation is not None: - # Broadcast the head now, during forward N, so - # followers build it during theirs (-1: not parallel). - speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( - speculation.node_batch, - speculative=True, spec_from_seq=pending.tp_seq, - ) - if self._tp_lead_needs_marker(pending, speculation): - self._broadcast_tp_nospec(pending) - # ``yield_away_from_target`` stays None when the admit - # refused, so the non-speculative path below schedules - # without the fairness exclusion — it should be free to pick - # up the batch ``_handle_admit_failure`` just pushed back. - if speculation is None and not admit_refused: - yield_away_from_target = ( - pending.node_name, - pending.graph_walk, - ) if must_yield_away else None - with self._span("worker.schedule_yield_away"): - batch = self.scheduler.get_next_batch( - self.request_state, - exclude_target=yield_away_from_target, - ) - if batch is not None: - node_batch = self._build_executing_batch(batch) - batch_partition = self.request_state.get_partition_for_node(batch.node_name) - logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") - speculation = Speculation( - scheduled_batch=batch, - node_batch=node_batch, - consumed_edges=set(), - continuing_rids=set(), # n/a - partition=batch_partition, - is_new_iter=False, - is_same_node=False, - is_yield_away=True - ) - - # A leader stamps the seq it sends; a follower's - # batch keeps the one it came off the FIFO with. - ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) - speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq - - def _arm_speculation(spec: Speculation) -> None: - # Hand N+1 to the plan thread NOW. It waits on N's - # commit event, then reserves a slot and pre-plans — - # while the main thread sits in await_gpu with the GIL - # released and N's kernels are still running. - if plan_executor is not None: - spec.plan_future = plan_executor.submit( - self._preplan_spec, pending, spec, - ) + if speculation is not None: + _arm_speculation(speculation) - if speculation is not None: - _arm_speculation(speculation) # 3. If pending: await GPU(N), submit speculated GPU(N+1) # asap, then post-process N (fast then slow) overlapping @@ -3921,11 +4021,42 @@ def _maybe_clear_spec(): # failed rids upstream. if pending.node_batch.admit_error is None: with self._span("worker.postprocess_batch"): - self._postprocess_batch(pending, outputs) - - # Removes for any rid not in the in-flight spec step - # are safe to apply now. - in_flight = set(spec_pending.batch.request_to_worker_graph) if spec_pending else set() + pp_state = self._postprocess_batch_stops(pending, outputs) + if ( + early_spec and pp_state is not None + and spec_pending is not None + and spec_pending.node_name not in self.parallel_nodes + ): + # N's stops are decided and N+1 is launched: + # build N+2 now and hand its pre-plan to the + # plan thread, so it overlaps the routing half + # below instead of the next launch. The fence + # first: the build moves pages, which must not + # happen while N+1's admit is outstanding. + with self._span("worker.speculate_early"): + self._await_admit_settled(spec_pending) + if spec_pending.node_batch.admit_error is None: + next_speculation, next_yield_away = ( + _build_speculation( + spec_pending, + consecutive_spec_steps + 1, + ) + ) + if next_speculation is not None: + self._armed_spec_rids = set( + next_speculation.scheduled_batch + .request_to_worker_graph + ) + self._in_flight_rids |= self._armed_spec_rids + _arm_speculation(next_speculation, spec_pending) + if pp_state is not None: + self._postprocess_batch_send(pending, outputs, pp_state) + + # Removes for any rid not in the in-flight spec step (or in + # a speculation armed ahead) are safe to apply now. + in_flight = ( + set(spec_pending.batch.request_to_worker_graph) if spec_pending else set() + ) | self._armed_spec_rids self._apply_pending_removes_safe_to_drop(in_flight) self._apply_pending_drains(in_flight) _set_pending(None) @@ -4001,7 +4132,11 @@ def _maybe_clear_spec(): phase_iter[0] += 1 _phase_flush() except Exception as e: - self._handle_main_loop_error(e, (pending, spec_pending), batch, speculation) + self._handle_main_loop_error( + e, (pending, spec_pending), batch, speculation, next_speculation, + ) + next_speculation = next_yield_away = None + self._armed_spec_rids = set() # Follower: a head from a step that raised must not sit at the # FIFO front with failed rids. if pending is not None and self._is_tp_follow_pending(pending): From 65ffc1aa2fa6d6a9dc37539b6e5f7a4ebedf5e0e Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 16:57:10 -0700 Subject: [PATCH 039/219] docs: MSTAR_EARLY_SPEC --- docs/environment_variables.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 878435e580..37018f301f 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -295,6 +295,15 @@ Worker scheduling - Pre-plan the speculative batch's attention on a dedicated thread while the previous replay runs. ``0`` plans inline on the GPU thread. + * - ``MSTAR_EARLY_SPEC`` + - ``1`` + - Build and pre-plan step N+2 while N+1 runs, between the stop and + routing halves of step N's post-processing, so the pre-plan overlaps + work the main thread does anyway instead of the next launch. ``0`` + builds each speculation at the top of the next iteration, after the + whole post-processing, which left the plan thread idle during it. + Non-parallel (TP1) nodes only; a tensor-parallel node keeps the old + order so the followers see heads in the order they settle them. * - ``MSTAR_LAZY_PUBLISH`` - ``1`` - Publish resource state (KV pages, positions) only for the requests From 7b21494a76eb2f85427bcd5750dde6c0692e5f29 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:24:08 -0700 Subject: [PATCH 040/219] sampler: skip the static gather for a steady batch while other requests wait for slot init --- mstar/engine/resources/sampler/utils.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index 7e6d82c193..d89149da3b 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1596,16 +1596,18 @@ def gather_static( update, never step to step. Skipped when ``cg_slot`` last gathered this same batch and no master row - changed since. Never skipped with a slot awaiting init, which also - catches a request id re-registered onto a new slot.""" + changed since. Never skipped while a row of this batch awaits init, + which also catches a request id re-registered onto a new slot. Slots + awaiting init that belong to requests outside the batch (registered, + still queued) do not block the skip: their rows are not read here, and + their init runs on the gather that first includes them.""" rids = tuple(request_ids) - if ( - not self._pending_init - and self._static_key.get(cg_slot) - == (rids, padded_bs, self._config_version) - ): - self._last_real_bs[cg_slot] = len(request_ids) - return + if self._static_key.get(cg_slot) == (rids, padded_bs, self._config_version): + pending = self._pending_init + get = self._rid_to_slot.get + if not pending or not any(get(rid) in pending for rid in request_ids): + self._last_real_bs[cg_slot] = len(request_ids) + return self._stage_slot_idx(request_ids, padded_bs, cg_slot) # H2D copies only (pre-plan, see HostBuffer); gather_dynamic uploads # the device-side index row From 6be4dd19de998e062b30ea6b8b5ccad22fc72804 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:24:08 -0700 Subject: [PATCH 041/219] kv plan: build the paged indptr tensors through numpy --- mstar/engine/resources/kv/plan.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/mstar/engine/resources/kv/plan.py b/mstar/engine/resources/kv/plan.py index d88c5c027c..c0be0d33a9 100644 --- a/mstar/engine/resources/kv/plan.py +++ b/mstar/engine/resources/kv/plan.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from typing import NamedTuple +import numpy as np import torch from mstar.engine.resources.step import Segment @@ -102,11 +103,12 @@ def build_paged_indptrs( segments: list[SequenceView], page_size: int, ) -> PagedIndptrs: - # TODO: ~90% of this is the four `torch.tensor(list)` conversions. Holding - # `CacheStream.page_indices` as a doubling int32 numpy array instead makes - # the pages copy a memcpy: measured 52us -> 13us at bs=16/1024 pages, for - # +0.3us on the per-step append. Not done because page_indices also crosses - # publish/ZMQ, offload/reload, forks and the CPU pool. + # Most of this is the five list-to-tensor conversions; they go through + # numpy, which builds an int32 array from a list of ints in about half + # the time `torch.tensor(list)` takes (42 -> 29 us at bs=32, 10 pages a + # row). Holding `CacheStream.page_indices` as a numpy array would cut the + # pages copy further, but it also crosses publish/ZMQ, offload/reload, + # forks and the CPU pool. qo_indptr = [0] kv_indptr = [0] all_pages: list[int] = [] @@ -122,14 +124,18 @@ def build_paged_indptrs( # FlashInfer's `get_seq_lens`, on the ints already here kv_lens.append(max(num_pages - 1, 0) * page_size + last) return PagedIndptrs( - qo_indptr=torch.tensor(qo_indptr, dtype=torch.int32), - paged_kv_indptr=torch.tensor(kv_indptr, dtype=torch.int32), - paged_kv_indices=torch.tensor(all_pages, dtype=torch.int32), - paged_kv_last_page_len=torch.tensor(last_page_lens, dtype=torch.int32), - kv_lens=torch.tensor(kv_lens, dtype=torch.int32), + qo_indptr=_int32_tensor(qo_indptr), + paged_kv_indptr=_int32_tensor(kv_indptr), + paged_kv_indices=_int32_tensor(all_pages), + paged_kv_last_page_len=_int32_tensor(last_page_lens), + kv_lens=_int32_tensor(kv_lens), ) +def _int32_tensor(values: list[int]) -> torch.Tensor: + return torch.from_numpy(np.array(values, dtype=np.int32)) + + @dataclass class KVPlanOutput: """ From ab3c97a924c8e6e58b3d40be0f94c84f055fb75d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:24:08 -0700 Subject: [PATCH 042/219] test: static gather skip with queued requests --- .../test_sampler_static_gather_skip.py | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 test/modular/test_sampler_static_gather_skip.py diff --git a/test/modular/test_sampler_static_gather_skip.py b/test/modular/test_sampler_static_gather_skip.py new file mode 100644 index 0000000000..11b748dd59 --- /dev/null +++ b/test/modular/test_sampler_static_gather_skip.py @@ -0,0 +1,74 @@ +"""``SamplerBuffers.gather_static`` skips a steady batch even while requests +outside the batch wait for their slot init (registered, still queued).""" +import pytest +import torch + +from mstar.engine.resources.sampler.utils import SamplerBuffers, SamplingConfig + + +def _buffers(monkeypatch) -> tuple[SamplerBuffers, list]: + buffers = SamplerBuffers.allocate(8, torch.device("cpu"), cg_slots=2) + gathers: list = [] + for buf in buffers._scalar_buffers(): + real = buf.gather + + def counting(rows, padded_bs, cg_slot, _real=real): + gathers.append(cg_slot) + return _real(rows, padded_bs, cg_slot) + + monkeypatch.setattr(buf, "gather", counting) + return buffers, gathers + + +def test_queued_request_does_not_block_the_skip(monkeypatch): + buffers, gathers = _buffers(monkeypatch) + buffers.register_request("a", SamplingConfig(temperature=0.5)) + buffers.register_request("b", SamplingConfig(temperature=0.7)) + buffers.gather_static(["a", "b"], 4, 0) + first = len(gathers) + assert first > 0 + + # a request registered but not yet in the batch: its slot awaits init + buffers.register_request("c", SamplingConfig(temperature=0.9)) + assert buffers._pending_init + buffers.gather_static(["a", "b"], 4, 0) + assert len(gathers) == first, "steady batch gathered again" + assert buffers._last_real_bs[0] == 2 + + # once it joins the batch, the gather runs and inits its row + buffers.gather_static(["a", "b", "c"], 4, 0) + assert len(gathers) > first + assert not buffers._pending_init + slot_c = buffers._rid_to_slot["c"] + assert float(buffers.temperature.master[slot_c]) == pytest.approx(0.9) + + +def test_pending_row_in_the_batch_forces_the_gather(monkeypatch): + buffers, gathers = _buffers(monkeypatch) + buffers.register_request("a", SamplingConfig(temperature=0.5)) + buffers.register_request("b", SamplingConfig(temperature=0.7)) + buffers.gather_static(["a", "b"], 4, 0) + first = len(gathers) + + # re-registered onto a new slot: the key still matches, the row does not + buffers.unregister_request("a") + buffers.register_request("a", SamplingConfig(temperature=0.2)) + buffers.gather_static(["a", "b"], 4, 0) + assert len(gathers) > first + slot_a = buffers._rid_to_slot["a"] + assert float(buffers.temperature.master[slot_a]) == pytest.approx(0.2) + + +def test_config_change_and_other_slot_still_gather(monkeypatch): + buffers, gathers = _buffers(monkeypatch) + buffers.register_request("a", SamplingConfig(temperature=0.5)) + buffers.gather_static(["a"], 2, 0) + first = len(gathers) + buffers.gather_static(["a"], 2, 1) # the other cg slot has nothing staged + second = len(gathers) + assert second > first + buffers.gather_static(["a"], 2, 1) + assert len(gathers) == second + buffers.update_request_config("a", SamplingConfig(temperature=0.1)) + buffers.gather_static(["a"], 2, 1) + assert len(gathers) > second From 1025fd5d01b5fba45c7b7c8f77e393d0494d28fb Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:24:08 -0700 Subject: [PATCH 043/219] ci: list the new cpu tests --- test/cpu-core.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index b5d8fec524..ab5e8d608f 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -44,3 +44,7 @@ test/modular/test_mamba2_resource.py test/modular/test_streams.py test/modular/test_batched_model_output.py test/modular/test_qwen3_5_check_stop_batched.py +test/modular/test_lazy_publish.py +test/modular/test_collect_outputs_fast_path.py +test/modular/test_inline_tokens.py +test/modular/test_sampler_static_gather_skip.py From 7359e679e191fd3f4f953a3a6d75475865521452 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 044/219] batched output: carry the row views the collect fast path built --- mstar/model/submodule_base.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index de9fc00f2f..d2d22d668d 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -80,6 +80,10 @@ class BatchedModelOutput: # the stop check, and slicing rows against that list would hand one # request another's token. None when there are no row-addressed buffers. row_request_ids: tuple[str, ...] | None = None + # Set by ``Engine._collect_outputs`` when every per-rid entry is a row view + # of a batch clone: name -> the views, row i for ``row_request_ids[i]``. + # The store describes such rows without reading each view. + row_views: dict[str, tuple[torch.Tensor, ...]] | None = None @classmethod def coerce(cls, output: BatchedModelOutput | dict[str, Any]) -> BatchedModelOutput: @@ -169,6 +173,8 @@ def update(self, other: BatchedModelOutput | dict[str, Any]): other = self.coerce(other) self.per_rid_outputs.update(other.per_rid_outputs) self.packed_outputs.update(other.packed_outputs) + # merged entries are no longer exactly one batch's row views + self.row_views = None # Row-addressed buffers describe one forward, so keep them only when # exactly one side had any; otherwise the stop check falls back to the # per-rid outputs. From 13e7943004ced720195ef2b9e7ba3d87285746b2 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 045/219] engine: hand the row views to the collected output --- mstar/engine/engine.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 5a9d63f45d..a68c437896 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1498,7 +1498,7 @@ def _collect_outputs( ) outputs: dict[str, NameToTensorList] = {} - row_clones = self._merge_per_rid( + row_clones, row_views = self._merge_per_rid( outputs, raw_outputs, request_ids, out_ids, submodule, req_info, ) self._merge_unpacked( @@ -1515,6 +1515,7 @@ def _collect_outputs( # row i is request_ids[i] as the forward ran them, not the # worker's later-rewritten request list row_request_ids=tuple(request_ids), + row_views=row_views, ) def _merge_per_rid( @@ -1525,13 +1526,15 @@ def _merge_per_rid( out_ids: list[str], submodule: NodeSubmodule, req_info: Mapping[str, CurrentForwardPassInfo], - ) -> dict[str, torch.Tensor]: + ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[torch.Tensor, ...]] | None]: """Fold the forward's per-rid entries into ``outputs``. ``row_outputs`` are cloned once for the batch and handed out as row views; per-rid entries are cloned individually after ``filter_batched_output``. A name in both resolves to the per-rid one. - Returns the row clones so the caller can share them. + Returns the row clones so the caller can share them, and on the + rows-only path the views themselves (name -> views, row i for + ``request_ids[i]``), which the store describes without reading each. """ row_clones = raw_outputs.clone_row_outputs(len(request_ids)) # The common decode case: only row outputs, and a submodule that does @@ -1555,7 +1558,7 @@ def _merge_per_rid( for name, views in row_views.items(): if i < len(views): merged[name] = [views[i]] - return row_clones + return row_clones, (row_views or None) for i, (rid, out_id) in enumerate( zip(request_ids, out_ids, strict=False) ): @@ -1588,7 +1591,7 @@ def _merge_per_rid( merged[key] = [value.clone()] else: merged[key] = value - return row_clones + return row_clones, None def _merge_unpacked( self, From 5a057b6c501a2c8008e3c5c5bd5560f97ad73406 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 046/219] tensor store: batch the ownership maps, dtype index helper --- mstar/communication/tensor_store.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/mstar/communication/tensor_store.py b/mstar/communication/tensor_store.py index af0392a288..f651983b98 100644 --- a/mstar/communication/tensor_store.py +++ b/mstar/communication/tensor_store.py @@ -69,6 +69,14 @@ class ColumnarTensorInfo: def __post_init__(self): self._seen_dtypes: dict[torch.dtype, int] = {} + def dtype_index(self, dtype: torch.dtype) -> int: + """The batch-local index of ``dtype``, registering it on first use.""" + idx = self._seen_dtypes.get(dtype) + if idx is None: + idx = self._seen_dtypes[dtype] = len(self.dtype_names) + self.dtype_names.append(_dtype_name(dtype)) + return idx + def add_tensors_canonical( self, uuids: list[int], canonicals: list[torch.Tensor], @@ -638,10 +646,15 @@ def _own( """Record which request owns each tensor. Pure Python dicts, so this stays a loop -- only the bookkeeping call crosses into Rust. ``cpu_tensors``, when given, is index-parallel to ``tensors``.""" - for request_id, uuid, tensor in zip(request_ids, uuids, tensors, strict=True): - self._tensors[uuid] = tensor - self._rid_to_uuids.setdefault(request_id, set()).add(uuid) - self._uuid_to_rid[uuid] = request_id + if len(request_ids) != len(uuids) or len(uuids) != len(tensors): + raise ValueError("request_ids, uuids and tensors must be parallel") + # two of the three maps take the batch in one update each; only the + # per-request set stays a loop + self._tensors.update(zip(uuids, tensors, strict=True)) + self._uuid_to_rid.update(zip(uuids, request_ids, strict=True)) + owned = self._rid_to_uuids + for request_id, uuid in zip(request_ids, uuids, strict=True): + owned.setdefault(request_id, set()).add(uuid) if cpu_tensors is not None: for uuid, cpu_tensor in zip(uuids, cpu_tensors, strict=True): if cpu_tensor is not None: From cb3d7e66bbd65ac64614294f4e47079e89d5b155 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 047/219] tensors: row store for a step whose outputs are row views of one batch clone --- mstar/communication/tensors.py | 173 +++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index 4ef7669bec..d9229a88d3 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -547,6 +547,179 @@ def store_and_return_tensor_info( ) return tensor_info + def store_row_outputs_batch( + self, + request_ids: list[Rid], + outputs, + signals: list[str], + row_views: dict[str, tuple[torch.Tensor, ...]], + row_request_ids: tuple[Rid, ...], + node_name: str | None = None, + graph_walk: str | None = None, + host_rows=None, + ) -> StoredOutputs | None: + """``store_and_return_tensor_info_batch`` for a step whose per-request + outputs are views of row-addressed batch tensors (a decode step: row i + of ``new_token`` is request i). + + The rows of one batch tensor differ only in address, so their + descriptors are built from one row's shape, stride and dtype instead + of being read off every view, and the host copies come from the stop + check's row buffers (``host_rows``) by index rather than being matched + tensor by tensor. Every (request, signal) entry must be exactly one + tensor that *is* a row view, or the result is None and the caller + takes the general path; nothing is recorded before that is settled. + What is stored is what the general path would store: one uuid per + (request, signal), request-major, the same row views. + """ + if not row_views or not self.tensor_store.has_put_tensor_batch_columns: + return None + row_of = {rid: i for i, rid in enumerate(row_request_ids)} + sharding = self.sharding_configs + first_cfg = sharding.get(request_ids[0]) if request_ids else None + tp_size, tp_rank = self._source_tp(first_cfg, node_name, graph_walk) + columns = ColumnarTensorInfo( + source_session_id=self.my_session_id, + source_entity=self.my_entity_id, + source_node_name=node_name, + source_graph_walk=graph_walk, + tp_size=tp_size, + tp_rank=tp_rank, + ) + want_cpu = host_rows is not None and self.needs_cpu_tensor + # signal -> the row tensor it views, learned from the first request + # that has it and checked on every other + alias: dict[str, str] = {} + # row tensor -> (base address, row step, nbytes, ndims, dims, stride, + # dtype index, host rows or None) + template: dict[str, tuple] = {} + flat_uuids: list[int] = [] + flat_rids: list[int] = [] + signal_idxs: list[int] = [] + num_tensors: list[int] = [] + canonicals: list[torch.Tensor] = [] + cpus: list[torch.Tensor | None] | None = [] if want_cpu else None + sharded_uuids: list[int] = [] + mint = self._uuid_minter.mint + add_uuid = columns.uuids.append + add_addr = columns.addresses.append + add_nbytes = columns.nbytes.append + add_ndims = columns.ndims.append + add_dims = columns.dims_flat.extend + add_stride = columns.stride_flat.extend + add_dtype = columns.dtype_idx.append + try: + for rid in request_ids: + i = row_of.get(rid) + if i is None: + return None + rid_out = outputs.get(rid) + cfg = sharding.get(rid) + if cfg is not first_cfg and self._source_tp( + cfg, node_name, graph_walk, + ) != (tp_size, tp_rank): + raise ValueError( + f"request {rid} has sharding " + f"{self._source_tp(cfg, node_name, graph_walk)} for " + f"node {node_name!r} walk {graph_walk!r}, but the batch " + f"was built for {(tp_size, tp_rank)}; one batch is one " + "node on one worker and cannot mix sharding groups" + ) + for s, signal in enumerate(signals): + lst = rid_out.get(signal) if rid_out else None + if not lst: + num_tensors.append(0) + continue + if type(lst) is not list or len(lst) != 1: + return None + tensor = lst[0] + name = alias.get(signal) + if name is None: + for cand, views in row_views.items(): + if views[i] is tensor: + name = cand + break + if name is None: + return None + alias[signal] = name + if name not in template: + template[name] = self._row_template( + views, columns, host_rows if want_cpu else None, name, + ) + elif row_views[name][i] is not tensor: + return None + # a sharded output is permuted on store; not a plain row + if cfg is not None and cfg.shard_dim.get(signal): + return None + base, step, nbytes, ndims, dims, stride, dtype_idx, host = template[name] + uuid = mint() + flat_uuids.append(uuid) + flat_rids.append(rid) + signal_idxs.append(s) + num_tensors.append(1) + canonicals.append(tensor) + add_uuid(uuid) + add_addr(base + i * step) + add_nbytes(nbytes) + add_ndims(ndims) + add_dims(dims) + add_stride(stride) + add_dtype(dtype_idx) + if want_cpu: + cpus.append(host[i] if host is not None else None) + if cfg is not None: + sharded_uuids.append(uuid) + except IndexError: + # a row index past the views or the host rows: not this batch's rows + return None + # nothing above touched shared state; from here the general path's + # guarantees apply + shard_dims = self.uuid_to_shard_dim + for uuid in sharded_uuids: + shard_dims[uuid] = None + edge_names = self.uuid_to_edge_name + if self.enable_prof: + for uuid, s in zip(flat_uuids, signal_idxs, strict=True): + edge_names[uuid] = signals[s] + try: + self.tensor_store.put_tensor_batch_columns( + flat_rids, canonicals, columns, cpu_tensors=cpus, + ) + except Exception: + for uuid in flat_uuids: + shard_dims.pop(uuid, None) + edge_names.pop(uuid, None) + raise + return StoredOutputs(flat_uuids, flat_rids, signal_idxs, num_tensors) + + @staticmethod + def _row_template( + views: tuple[torch.Tensor, ...], columns: ColumnarTensorInfo, + host_rows, name: str, + ) -> tuple: + """What every row of one batch tensor shares, for the row store: the + first row's address and the step to the next, byte size, shape, + stride and dtype index, plus the matching host rows (one view per row) + when ``host_rows`` carries a buffer of this name with the rows' shape + and dtype.""" + first = views[0] + base = first.data_ptr() + step = views[1].data_ptr() - base if len(views) > 1 else 0 + dims = tuple(first.shape) + host = None + if host_rows is not None: + buf = host_rows.buffers.get(name) + if ( + torch.is_tensor(buf) and buf.device.type == "cpu" + and buf.dim() == len(dims) and buf.dtype == first.dtype + and tuple(buf.shape[1:]) == dims[1:] and buf.shape[0] >= len(views) + ): + host = buf.split(1) + return ( + base, step, first.nbytes, len(dims), dims, tuple(first.stride()), + columns.dtype_index(first.dtype), host, + ) + def store_and_return_tensor_info_batch( self, request_ids: list[Rid], From f19248a13c002a29fba86c9f7c823f38059fbe04 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 048/219] worker: store decode row outputs through the row store --- mstar/worker/worker.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 70f52dd3a7..f84e1f2683 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2909,13 +2909,24 @@ def _pp_stage(name: str) -> None: # Pass the stop check's host copies so a host-memory transport needn't # copy the rows again. They are views of pinned buffers the next step # reuses; safe because the sends below are their last reader. - stored = self.tensor_manager.store_and_return_tensor_info_batch( - rids, outputs, signals, - node_name=batch_N.node_name, - graph_walk=batch_N.graph_walk, - skip_cuda_sync=True, - cpu_tensors=cpu_outputs, - ) + # Row views of one batch clone (a decode step) are described from one + # row; anything else goes through the general per-tensor path. + stored = None + if outputs.row_views and outputs.row_request_ids is not None: + stored = self.tensor_manager.store_row_outputs_batch( + rids, outputs, signals, outputs.row_views, outputs.row_request_ids, + node_name=batch_N.node_name, + graph_walk=batch_N.graph_walk, + host_rows=host_rows, + ) + if stored is None: + stored = self.tensor_manager.store_and_return_tensor_info_batch( + rids, outputs, signals, + node_name=batch_N.node_name, + graph_walk=batch_N.graph_walk, + skip_cuda_sync=True, + cpu_tensors=cpu_outputs, + ) flat_uuids = stored.flat_uuids flat_rids = stored.flat_rids signal_idxs = stored.signal_idxs From 129311f90081729ce8cb6f57700cc608efbe70f7 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 049/219] test: merge_per_rid returns the row views --- test/modular/test_collect_outputs_fast_path.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_collect_outputs_fast_path.py b/test/modular/test_collect_outputs_fast_path.py index 4db4db772d..db3a049aa6 100644 --- a/test/modular/test_collect_outputs_fast_path.py +++ b/test/modular/test_collect_outputs_fast_path.py @@ -44,7 +44,7 @@ def _raw(bs=4, with_per_rid=False): def _collect(sub, raw, rids): eng = Engine.__new__(Engine) outputs = {} - row_clones = eng._merge_per_rid(outputs, raw, rids, rids, sub, {}) + row_clones, _ = eng._merge_per_rid(outputs, raw, rids, rids, sub, {}) return outputs, row_clones From c104c997913edc75d1195f17102ce672936aed5b Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 050/219] test: row store against the general path --- test/modular/test_row_output_store.py | 121 ++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 test/modular/test_row_output_store.py diff --git a/test/modular/test_row_output_store.py b/test/modular/test_row_output_store.py new file mode 100644 index 0000000000..78c9d814ae --- /dev/null +++ b/test/modular/test_row_output_store.py @@ -0,0 +1,121 @@ +"""The row store: a decode step's outputs are row views of one batch clone, +stored from one row's description and the stop check's host rows.""" +import torch + +from mstar.communication.tensors import SharedMemoryCommunicationManager +from mstar.model.submodule_base import BatchedModelOutput, HostRows + +SIGNALS = ["text_inputs", "new_token"] + + +def _manager(tmp_path): + class _NullCommunicator: + def send(self, *args, **kwargs): + pass + + return SharedMemoryCommunicationManager( + my_entity_id="worker_0", hostname="localhost", device="cpu", + communicator=_NullCommunicator(), shm_dir=str(tmp_path), + ) + + +def _batch(rids): + rows = torch.arange(10, 10 + len(rids), dtype=torch.long).view(len(rids), 1) + clone = rows.clone() + views = clone.split(1) + per_rid = {} + for i, rid in enumerate(rids): + per_rid[rid] = {"new_token": [views[i]]} + # the model's postprocess rebinds the loop-back name to the same list + per_rid[rid]["text_inputs"] = per_rid[rid]["new_token"] + outputs = BatchedModelOutput( + per_rid_outputs=per_rid, row_views={"new_token": views}, + row_request_ids=tuple(rids), + ) + host = rows.clone() + host_rows = HostRows(request_ids=tuple(rids), buffers={"new_token": host}) + cpu_per_rid = { + rid: {"new_token": [host[i : i + 1]], "text_inputs": [host[i : i + 1]]} + for i, rid in enumerate(rids) + } + return outputs, host_rows, cpu_per_rid + + +def test_row_store_matches_the_general_path(tmp_path): + rids = [3, 5, 8, 13] + outputs, host_rows, cpu_per_rid = _batch(rids) + fast_mgr, slow_mgr = _manager(tmp_path / "a"), _manager(tmp_path / "b") + fast = fast_mgr.store_row_outputs_batch( + rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, + node_name="text", graph_walk="decode", host_rows=host_rows, + ) + slow = slow_mgr.store_and_return_tensor_info_batch( + rids, outputs, SIGNALS, node_name="text", graph_walk="decode", + cpu_tensors=cpu_per_rid, + ) + assert fast is not None + assert fast.flat_rids == slow.flat_rids == [3, 3, 5, 5, 8, 8, 13, 13] + assert fast.signal_idxs == slow.signal_idxs == [0, 1] * 4 + assert fast.num_tensors == slow.num_tensors == [1] * 8 + assert len(set(fast.flat_uuids)) == 8 + for uf, us in zip(fast.flat_uuids, slow.flat_uuids, strict=True): + tf, ts = fast_mgr.get_tensor(uf), slow_mgr.get_tensor(us) + assert tf.data_ptr() == ts.data_ptr() and tf.shape == ts.shape + assert tf.dtype == ts.dtype and tf.stride() == ts.stride() + info_f, info_s = fast_mgr.tensor_store.get_info(uf), slow_mgr.tensor_store.get_info(us) + assert (info_f.dims, info_f.stride, info_f.nbytes, info_f.dtype, info_f.address) == ( + info_s.dims, info_s.stride, info_s.nbytes, info_s.dtype, info_s.address + ) + cf = fast_mgr.tensor_store._tensors_cpu.get(uf) + cs = slow_mgr.tensor_store._tensors_cpu.get(us) + assert cf is not None and cs is not None + assert torch.equal(cf, cs) and cf.data_ptr() == cs.data_ptr() + for rid in rids: + assert len(fast_mgr.tensor_store.get_all_uuids(rid)) == 2 + + +def test_row_store_declines_what_is_not_a_row_view(tmp_path): + rids = [1, 2, 3] + outputs, host_rows, _ = _batch(rids) + mgr = _manager(tmp_path) + # one request's output was cloned per request: not a row of the batch + outputs.per_rid_outputs[2]["new_token"] = [outputs.per_rid_outputs[2]["new_token"][0].clone()] + outputs.per_rid_outputs[2]["text_inputs"] = outputs.per_rid_outputs[2]["new_token"] + assert mgr.store_row_outputs_batch( + rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, + host_rows=host_rows, + ) is None + assert mgr.tensor_store.get_all_uuids(1) == [] + + # a request that has no row at all + outputs2, host_rows2, _ = _batch(rids) + assert mgr.store_row_outputs_batch( + rids + [9], outputs2, SIGNALS, outputs2.row_views, outputs2.row_request_ids, + host_rows=host_rows2, + ) is None + + +def test_row_store_keeps_the_general_path_shape_for_missing_signals(tmp_path): + rids = [1, 2, 3] + outputs, host_rows, cpu_per_rid = _batch(rids) + del outputs.per_rid_outputs[2]["text_inputs"] + mgr = _manager(tmp_path) + fast = mgr.store_row_outputs_batch( + rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, + host_rows=host_rows, + ) + assert fast is not None + assert fast.num_tensors == [1, 1, 0, 1, 1, 1] + assert fast.flat_rids == [1, 1, 2, 3, 3] + assert fast.signal_idxs == [0, 1, 1, 0, 1] + + +def test_row_store_without_host_rows_stores_no_host_copy(tmp_path): + rids = [1, 2] + outputs, _, _ = _batch(rids) + mgr = _manager(tmp_path) + fast = mgr.store_row_outputs_batch( + rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, + ) + assert fast is not None and len(fast.flat_uuids) == 4 + assert not mgr.tensor_store._tensors_cpu From 28d739d9638afb651cfe4abd199f339d4947ee38 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:55:24 -0700 Subject: [PATCH 051/219] ci: list the row store test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index ab5e8d608f..5d3711984a 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -48,3 +48,4 @@ test/modular/test_lazy_publish.py test/modular/test_collect_outputs_fast_path.py test/modular/test_inline_tokens.py test/modular/test_sampler_static_gather_skip.py +test/modular/test_row_output_store.py From 760bde478b4c9dd1ee1722eaafff186b5055842f Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:58:35 -0700 Subject: [PATCH 052/219] worker: build the stop check's per-rid host dicts only when read --- mstar/worker/worker.py | 48 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index f84e1f2683..51bb3d05f9 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -221,6 +221,52 @@ class _PostprocessState: signal_idxs: object = None +class _LazyPerRid: + """``Worker._rows_to_per_rid`` on demand: the per-rid host dicts are built + the first time a consumer reads them. The batched stop check, the row store + and the inline emit read the rows directly, so in the common decode step + nobody does, and the per-request dicts are never built.""" + + __slots__ = ("_host", "_request_ids", "_built") + + def __init__(self, host: dict, request_ids: list[int]): + self._host = host + self._request_ids = request_ids + self._built: dict | None = None + + def _dict(self) -> dict: + if self._built is None: + self._built = Worker._rows_to_per_rid(self._host, self._request_ids) + return self._built + + def get(self, key, default=None): + return self._dict().get(key, default) + + def __getitem__(self, key): + return self._dict()[key] + + def __contains__(self, key) -> bool: + return key in self._dict() + + def __iter__(self): + return iter(self._dict()) + + def __len__(self) -> int: + return len(self._request_ids) + + def __bool__(self) -> bool: + return bool(self._request_ids) + + def items(self): + return self._dict().items() + + def keys(self): + return self._dict().keys() + + def values(self): + return self._dict().values() + + class Worker: """ Real worker that integrates RequestStateManager, EngineManager, @@ -3280,7 +3326,7 @@ def _prematerialize_for_check_stop( if outputs.check_stop_buffers is not None and row_rids is not None: host = self._d2h_batched(outputs.check_stop_buffers, side) return ( - Worker._rows_to_per_rid(host, row_rids), + _LazyPerRid(host, row_rids), HostRows(tuple(row_rids), host), ) From 4e50fb79e7708f3724f3670d721af7135df6311a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:58:35 -0700 Subject: [PATCH 053/219] engine: skip the unpack pass when nothing is packed --- mstar/engine/engine.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index a68c437896..3c4944ef7b 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1607,6 +1607,13 @@ def _merge_unpacked( A captured region can't do this itself: the per-request slice ends depend on the real seq_lens, which only reach it through the plan. """ + if ( + not raw_outputs.packed_outputs + and type(submodule).unpack_packed_outputs + is NodeSubmodule.unpack_packed_outputs + ): + # nothing packed and the base no-op: skip the per-rid lengths + return unpacked = submodule.unpack_packed_outputs( static_output=raw_outputs.packed_outputs, request_ids=request_ids, From 91632550227998fae8a7fca469d6d10388c532a2 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:58:35 -0700 Subject: [PATCH 054/219] test: lazy per-rid host dicts --- test/modular/test_lazy_per_rid.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 test/modular/test_lazy_per_rid.py diff --git a/test/modular/test_lazy_per_rid.py b/test/modular/test_lazy_per_rid.py new file mode 100644 index 0000000000..f656e595dd --- /dev/null +++ b/test/modular/test_lazy_per_rid.py @@ -0,0 +1,22 @@ +"""The stop check's per-rid host dicts are built only when something reads +them.""" +import torch + +from mstar.worker.worker import Worker, _LazyPerRid + + +def test_lazy_per_rid_builds_on_first_read_only(): + host = {"new_token": torch.arange(3).view(3, 1), "flag": torch.zeros(3, dtype=torch.bool)} + rids = [7, 8, 9] + lazy = _LazyPerRid(host, rids) + assert lazy._built is None + assert len(lazy) == 3 and bool(lazy) + assert lazy._built is None, "len and truth must not build the dicts" + eager = Worker._rows_to_per_rid(host, rids) + assert lazy.get(8).keys() == eager[8].keys() + for name in eager[8]: + assert torch.equal(lazy.get(8)[name][0], eager[8][name][0]) + assert lazy.get(42) is None + assert 7 in lazy and list(lazy) == rids + assert lazy[9]["new_token"][0].item() == 2 + assert not _LazyPerRid(host, []) From 5e2b57c310e3741c246b2a9b7861fa5f8a457de3 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 17:58:35 -0700 Subject: [PATCH 055/219] ci: list the lazy per-rid test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 5d3711984a..e52b0e0bd6 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -49,3 +49,4 @@ test/modular/test_collect_outputs_fast_path.py test/modular/test_inline_tokens.py test/modular/test_sampler_static_gather_skip.py test/modular/test_row_output_store.py +test/modular/test_lazy_per_rid.py From 083cc0f5f590b005a62b46f26553fc179eab33c5 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:07:10 -0700 Subject: [PATCH 056/219] qwen3_5: MSTAR_QWEN35_DECODE_BUCKETS overrides the decode capture ladder --- mstar/model/qwen3_5/submodules.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index a29ce50e21..7b2dbaf006 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -7,6 +7,7 @@ """ import logging +import os from collections.abc import Mapping from typing import Any @@ -83,7 +84,13 @@ class LLMSubmodule(ARNodeSubmodule): # cap on rows per decode step: with 32 as the top, 64 requests in flight # ran as two alternating batches of 32 and each paid the whole per-step # host cost. A step above 32 rows pads to the next bucket. - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 48, 64, 96, 128] + # MSTAR_QWEN35_DECODE_BUCKETS="1,2,4,8,16,32" narrows the ladder (less graph + # memory, or the pre-#258 ladder for an A/B). + DECODE_CAPTURE_BATCH_SIZES = [ + int(x) for x in os.environ.get( + "MSTAR_QWEN35_DECODE_BUCKETS", "1,2,4,8,16,32,48,64,96,128", + ).split(",") + ] # Text plus image tokens of a whole prompt. Capped at 4096, not the # processor's 16384, because static buffers are sized by the largest From 5025706e8d8027af4e2abfbcdd1f0c179c8eb871 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:07:10 -0700 Subject: [PATCH 057/219] docs: MSTAR_QWEN35_DECODE_BUCKETS --- docs/environment_variables.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 37018f301f..5a50265888 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -295,6 +295,13 @@ Worker scheduling - Pre-plan the speculative batch's attention on a dedicated thread while the previous replay runs. ``0`` plans inline on the GPU thread. + * - ``MSTAR_QWEN35_DECODE_BUCKETS`` + - ``1,2,4,8,16,32,48,64,96,128`` + - The decode batch sizes Qwen3.5 captures as CUDA graphs. A decode step + runs at the smallest bucket that holds its rows, so the ladder bounds + both the rows per step and the graph memory; narrow it (for example + ``1,2,4,8,16,32``) on a memory-tight GPU or for an A/B against the + former ladder. * - ``MSTAR_EARLY_SPEC`` - ``1`` - Build and pre-plan step N+2 while N+1 runs, between the stop and From a025faca67d9f66daaf991be8c3e813e97d55556 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:19 -0700 Subject: [PATCH 058/219] step context: constant-time padding and real-row tests --- mstar/engine/resources/step.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/step.py b/mstar/engine/resources/step.py index cee3fda9f9..f746d22ca0 100644 --- a/mstar/engine/resources/step.py +++ b/mstar/engine/resources/step.py @@ -71,6 +71,10 @@ class StepContext: # None outside a captured replay, where the two are the same. Padding rows # carry negative handles, so this stays homogeneous with `request_ids`. _padded_request_ids: Sequence[int] | None = None + # ``request_ids`` as a set, built on first use after the padded list is + # set: the per-row tests below run for every row of every resource on two + # threads, and a list scan made each of them linear in the batch. + _real_rids: frozenset | None = field(default=None, repr=False) @property def padded_request_ids(self) -> Sequence[int]: @@ -80,13 +84,26 @@ def padded_request_ids(self) -> Sequence[int]: def set_padded_rids(self, padded_rids: Sequence[int] | None): self._padded_request_ids = padded_rids + self._real_rids = None + + def is_real_row(self, rid: str) -> bool: + """Whether ``rid`` is one of this step's real requests.""" + real = self._real_rids + if real is None: + real = self._real_rids = frozenset(self.request_ids) + return rid in real def is_padding_row(self, rid: str) -> bool: """Whether ``rid`` is a replay's padding row rather than a real request. It carries the template span so shapes match the capture, but a resource should reserve, address and commit nothing on its behalf.""" - return self._padded_request_ids is not None and rid not in self.request_ids + if self._padded_request_ids is None: + return False + real = self._real_rids + if real is None: + real = self._real_rids = frozenset(self.request_ids) + return rid not in real def set_piecewise_leases(self, leases: "Mapping[str, SlotLease]"): self.piecewise_leases = leases From b3a01b6ce603ec04f20ae24f67d9f4aa0e759504 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:19 -0700 Subject: [PATCH 059/219] kv: batched publish snapshot under one lock, real-row test in the page index --- mstar/engine/resources/kv/manager.py | 38 +++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 318629a68b..ab0d31515b 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -535,7 +535,7 @@ def _index_filled_pages( if ( self._index is None or ctx.capture - or segment.request_id not in ctx.request_ids + or not ctx.is_real_row(segment.request_id) or chain is None ): return @@ -1776,6 +1776,42 @@ def publish_for_step( request_id, node_name=node_name, graph_walk=graph_walk, ) + def publish_snapshot_batch( + self, + request_ids: list[str], + node_name: str | None, + graph_walk: str | None, + ) -> list: + """``publish_snapshot_for_step`` for a batch: one lock acquisition and + a plain loop, which is what the gpu thread can afford per row.""" + out = [] + append = out.append + streams_of = self._streams.get + overrides_of = self._overrides.get + with self._lock: + for request_id in request_ids: + streams = streams_of(request_id) + overrides = overrides_of(request_id) + if streams is None or overrides is None: + append(None) + continue + labels = overrides.get_publish_labels( + node_name, graph_walk, list(streams), final=False, + ) + if not labels: + append(None) + continue + snap = [] + for label in labels: + stream = streams.get(label) + if stream is not None: + snap.append(( + label, stream.stored_len, len(stream.page_indices), + stream.reset_generation, + )) + append(tuple(snap) if snap else None) + return out + def publish_snapshot_for_step( self, request_id: str, From 91e4d28ab7df37445f02267dc228330b5b056c1b Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:19 -0700 Subject: [PATCH 060/219] resources: publish_snapshot_batch hook --- mstar/engine/resources/base.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index 6ea4947740..c304ad4363 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -220,6 +220,22 @@ def publish_after_stop( del request_id, node_name, graph_walk return None + def publish_snapshot_batch( + self, + request_ids: list[str], + node_name: str | None, + graph_walk: str | None, + ) -> list: + """``publish_snapshot_for_step`` for every request of a step, one + entry per request in order (None for nothing to publish). A resource + with per-request locking or lookups overrides this to do them once.""" + return [ + self.publish_snapshot_for_step( + rid, node_name=node_name, graph_walk=graph_walk, + ) + for rid in request_ids + ] + def publish_snapshot_for_step( self, request_id: str, From b2f60e6335551d6b065058cbc7845c4c842e9296 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:19 -0700 Subject: [PATCH 061/219] runner: snapshot the publication per resource for the whole batch --- mstar/engine/resources/runner.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index 95f947cc02..14d54435e9 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -495,14 +495,19 @@ def publish_snapshot( order = self._sweep(self._node_publish_order, self._publish_order, node_name) if not order: return {} - publishers = [(key, self._resources[key]) for key in order] + # one call per publisher for the whole batch, then one dict per + # request that has anything to publish + columns = [ + (key, self._resources[key].publish_snapshot_batch( + request_ids, node_name=node_name, graph_walk=graph_walk, + )) + for key in order + ] out: dict[str, dict[str, Any]] = {} - for rid in request_ids: + for i, rid in enumerate(request_ids): per_key: dict[str, Any] = {} - for key, resource in publishers: - snap = resource.publish_snapshot_for_step( - rid, node_name=node_name, graph_walk=graph_walk, - ) + for key, snaps in columns: + snap = snaps[i] if snap is not None: per_key[key] = snap if per_key: From cb6bed178ce5421ad8fcfcc439b02e61251bb2a8 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:19 -0700 Subject: [PATCH 062/219] sampler: one graph sampler per bucket and slot --- mstar/engine/resources/sampler/utils.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index d89149da3b..c17af7a48d 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1327,6 +1327,8 @@ class SamplerBuffers: _static_key: dict[int, tuple[tuple[str, ...], int, int]] = field( default_factory=dict, repr=False ) + # (padded_bs, cg slot) -> the sampler over that slot's views; see sampler_for + _samplers: dict = field(default_factory=dict, repr=False) @property def tracks_seen_tokens(self) -> bool: @@ -1647,8 +1649,18 @@ def gather_dynamic( def sampler_for(self, padded_bs: int, cg_slot: int) -> "CudaGraphableSampler": """A sampler bound to ``cg_slot``'s per-step buffer views (zero-copy). - Valid regardless of when the buffers are (re)gathered into.""" - return CudaGraphableSampler(**self.slice_for_bs(padded_bs, cg_slot)) + Valid regardless of when the buffers are (re)gathered into. + + Cached per (bucket, slot): the per-step buffers are allocated once and + never rebound, so the views are the same every step, and the sampler + itself keeps no per-step state the next step could read.""" + key = (padded_bs, cg_slot) + sampler = self._samplers.get(key) + if sampler is None: + sampler = CudaGraphableSampler(**self.slice_for_bs(padded_bs, cg_slot)) + self._samplers[key] = sampler + sampler.applied_penalty_in_graph = False + return sampler def scatter_offset(self, cg_slot: int = 0) -> None: """Persist the (in-graph advanced) per-step offsets back to their slot From 1dce205dd9f837cf8eab77a20098875bf457753a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:20 -0700 Subject: [PATCH 063/219] test: step context row membership --- test/modular/test_step_context_rows.py | 29 ++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 test/modular/test_step_context_rows.py diff --git a/test/modular/test_step_context_rows.py b/test/modular/test_step_context_rows.py new file mode 100644 index 0000000000..990b8e2796 --- /dev/null +++ b/test/modular/test_step_context_rows.py @@ -0,0 +1,29 @@ +"""StepContext row tests: constant-time membership that follows the request +list through padding and rewrites.""" +from mstar.engine.resources.step import StepContext + + +def _ctx(rids): + return StepContext(request_ids=list(rids), graph_walk="decode", slot=0, capture=False) + + +def test_padding_rows_are_the_padded_tail(): + ctx = _ctx([5, 7]) + assert not ctx.is_padding_row(5) and not ctx.is_padding_row(99) + assert ctx.is_real_row(7) and not ctx.is_real_row(99) + ctx.set_padded_rids([5, 7, -1, -2]) + assert [ctx.is_padding_row(r) for r in (5, 7, -1, -2)] == [False, False, True, True] + assert ctx.padded_request_ids == [5, 7, -1, -2] + + +def test_rewriting_the_request_list_refreshes_the_set(): + ctx = _ctx([1, 2, 3]) + ctx.set_padded_rids([1, 2, 3, -1]) + assert ctx.is_padding_row(-1) and not ctx.is_padding_row(2) + # a dropped request: the batch rewrites the list and clears the padding + ctx.set_padded_rids(None) + ctx.request_ids = [1, 3] + assert not ctx.is_padding_row(2) # no padded list: nothing is padding + assert not ctx.is_real_row(2) and ctx.is_real_row(3) + ctx.set_padded_rids([1, 3, -1, -2]) + assert ctx.is_padding_row(2) and ctx.is_padding_row(-2) and not ctx.is_padding_row(3) From 6cb5601589f191a77c1d19f213fa6c6c197a6bc7 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:20 -0700 Subject: [PATCH 064/219] test: batched publish snapshot --- test/modular/test_publish_snapshot_batch.py | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 test/modular/test_publish_snapshot_batch.py diff --git a/test/modular/test_publish_snapshot_batch.py b/test/modular/test_publish_snapshot_batch.py new file mode 100644 index 0000000000..cdfc6a624f --- /dev/null +++ b/test/modular/test_publish_snapshot_batch.py @@ -0,0 +1,45 @@ +"""The runner snapshots a step's publication per resource for the whole batch, +and the result matches the per-request form.""" +from typing import Any + +from mstar.engine.resources.base import Resource +from mstar.engine.resources.runner import StepRunner + + +class _Pub: + """A publisher standing in for a resource: snapshots are per-rid tuples.""" + + def __init__(self, name, skip=()): + self.name = name + self.skip = set(skip) + self.batch_calls = 0 + + def publish_snapshot_for_step(self, request_id, node_name, graph_walk): + if request_id in self.skip: + return None + return (self.name, request_id, node_name, graph_walk) + + def publish_snapshot_batch(self, request_ids, node_name, graph_walk): + self.batch_calls += 1 + return Resource.publish_snapshot_batch(self, request_ids, node_name, graph_walk) + + +def _runner(pubs: dict[str, Any]) -> StepRunner: + runner = StepRunner.__new__(StepRunner) + runner._resources = pubs + runner._publish_order = list(pubs) + runner._node_publish_order = {} + runner._sweep = lambda node_order, order, node_name: order + return runner + + +def test_batch_snapshot_matches_per_request_form(): + pubs = {"kv": _Pub("kv", skip={3}), "pos": _Pub("pos", skip={2, 3})} + runner = _runner(pubs) + out = runner.publish_snapshot([1, 2, 3], node_name="LLM", graph_walk="decode") + assert out == { + 1: {"kv": ("kv", 1, "LLM", "decode"), "pos": ("pos", 1, "LLM", "decode")}, + 2: {"kv": ("kv", 2, "LLM", "decode")}, + } + assert pubs["kv"].batch_calls == 1 and pubs["pos"].batch_calls == 1 + assert runner.publish_snapshot([], node_name="LLM", graph_walk="decode") == {} From d2dcb3898417c9dd8fb373286f38bd356520eabe Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:20 -0700 Subject: [PATCH 065/219] test: sampler_for cache --- test/modular/test_sampler_for_cache.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 test/modular/test_sampler_for_cache.py diff --git a/test/modular/test_sampler_for_cache.py b/test/modular/test_sampler_for_cache.py new file mode 100644 index 0000000000..cad60545eb --- /dev/null +++ b/test/modular/test_sampler_for_cache.py @@ -0,0 +1,19 @@ +"""``SamplerBuffers.sampler_for`` hands back one sampler per (bucket, slot) +over the same buffer views.""" +import torch + +from mstar.engine.resources.sampler.utils import SamplerBuffers + + +def test_sampler_for_is_cached_per_bucket_and_slot(): + buffers = SamplerBuffers.allocate(8, torch.device("cpu"), cg_slots=2) + a = buffers.sampler_for(4, 0) + assert buffers.sampler_for(4, 0) is a + b = buffers.sampler_for(4, 1) + c = buffers.sampler_for(8, 0) + assert b is not a and c is not a and c is not b + assert a.temperature_buf.data_ptr() == buffers.temperature.slot_view(0, 4).data_ptr() + assert b.temperature_buf.data_ptr() == buffers.temperature.slot_view(1, 4).data_ptr() + assert a.temperature_buf.shape[0] == 4 and c.temperature_buf.shape[0] == 8 + a.applied_penalty_in_graph = True + assert buffers.sampler_for(4, 0).applied_penalty_in_graph is False From 9f222e4486fbcf12fc58930a31efa66a45c8bcd0 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:13:20 -0700 Subject: [PATCH 066/219] ci: list the step, snapshot and sampler tests --- test/cpu-core.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index e52b0e0bd6..52e9bb4408 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -50,3 +50,6 @@ test/modular/test_inline_tokens.py test/modular/test_sampler_static_gather_skip.py test/modular/test_row_output_store.py test/modular/test_lazy_per_rid.py +test/modular/test_step_context_rows.py +test/modular/test_publish_snapshot_batch.py +test/modular/test_sampler_for_cache.py From fb047710455722cf5c25db096adf6798be613db5 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:19:51 -0700 Subject: [PATCH 067/219] kv admit: skip the allocation call for rows whose pages already cover the step --- mstar/engine/resources/kv/manager.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index ab0d31515b..42e4d2188b 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -869,7 +869,10 @@ def admit(self, step: KVStep, ctx: StepContext) -> AdmitOutcome: # moved — re-admitting it allocates nothing and re-copies nothing. # A post-fork copies the source *after* this step's spans land, so its # reservation covers them. - growth = self._label_growth(step) if step.commit else {} + # only a post-fork reads the growth, so a decode step does not build it + growth = ( + self._label_growth(step) if step.commit and step.post_forks else {} + ) forks = [(pre, 0) for pre in step.pre_forks] + [ (post, growth) for post in step.post_forks ] @@ -927,6 +930,7 @@ def admit(self, step: KVStep, ctx: StepContext) -> AdmitOutcome: self._admit_reserved_pages.setdefault( (rid, to_label), []).extend(alloc_res.new_pages) + page_size = self.config.page_size for segment in step.segments: # A replay's padding rows reserve nothing: they run against # SINK_PAGE (see `_sequence_views`). Pages for them would fail @@ -934,10 +938,16 @@ def admit(self, step: KVStep, ctx: StepContext) -> AdmitOutcome: if segment.span == 0 or ctx.is_padding_row(segment.request_id): continue stream = self._ensure_label(segment.request_id, segment.label) + seq_len = segment.span + stream.stored_len + # the common decode row: its pages already cover the token it + # appends, so there is nothing to acquire and nothing to record + if ( + not stream.offloaded + and len(stream.page_indices) * page_size >= seq_len + ): + continue alloc_res = self._alloc( - segment.request_id, - segment.label, - segment.span + stream.stored_len + segment.request_id, segment.label, seq_len, ) if not alloc_res.success: return AdmitOutcome(ok=False, reason=alloc_res.error) From a2d1155f28a26551143171a2274b4265a18cb70d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:19:51 -0700 Subject: [PATCH 068/219] engine: hoist the per-row constants out of prepare_inputs, look up the prefix cache only on keyed walks --- mstar/engine/engine.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 3c4944ef7b..450b1e1010 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -716,22 +716,33 @@ def prepare_inputs(self, batch: ExecutingBatch) -> None: range_pop() def _prepare_inputs(self, batch: ExecutingBatch) -> None: - submodule = self._submodules[batch.node_name].submodule + mgmt = self._submodules[batch.node_name] + submodule = mgmt.submodule node_inputs: list[NodeInputs] = [] + # hoisted out of the per-request loop: the same for every row + walk = batch.step_context.graph_walk + prepare = submodule.prepare_inputs + resources = mgmt.resources + wrapped = batch.per_request_info_wrapped + tensors = batch.per_request_input_tensors + metadata = batch.per_request_input_metadata + final_rids = batch.final_stream_rids + # only a keyed walk probes the prefix cache (see _skip_cached_prefix) + probe_prefix = walk in self._keyed_walks.get(batch.node_name, ()) for rid in batch.request_ids: try: - req_inputs = submodule.prepare_inputs( - graph_walk=batch.step_context.graph_walk, - fwd_info=batch.per_request_info_wrapped[rid], - inputs=batch.per_request_input_tensors.get(rid, {}), - resources=self._submodules[batch.node_name].resources, + req_inputs = prepare( + graph_walk=walk, + fwd_info=wrapped[rid], + inputs=tensors.get(rid, {}), + resources=resources, # the step that ends the last of the node's streams: it must # flush whatever it held back (a vocoder's crossfade tail, the # look-ahead frames a token encoder withholds) - is_final_stream_chunk=rid in batch.final_stream_rids, - input_metadata=batch.per_request_input_metadata.get(rid, EMPTY_INPUT_METADATA), + is_final_stream_chunk=rid in final_rids, + input_metadata=metadata.get(rid, EMPTY_INPUT_METADATA), ) - if req_inputs is not None: + if req_inputs is not None and probe_prefix: req_inputs = self._skip_cached_prefix(batch, rid, req_inputs) except Exception as error: logger.exception( From dbbc572e8ba1e12f8387edc1c07a6832004f313d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:19:51 -0700 Subject: [PATCH 069/219] arena: batched cleanup of collectable tensors --- mstar/communication/arena.py | 41 ++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/mstar/communication/arena.py b/mstar/communication/arena.py index 1457e08a80..0dae61ce6b 100644 --- a/mstar/communication/arena.py +++ b/mstar/communication/arena.py @@ -925,6 +925,47 @@ def close(self) -> None: self._wake_q.put(None) self._wake_q = None + def cleanup_collectable( + self, collectable: list[int], registered: list[bool], + ): + """``_cleanup_by_uuid`` for a batch, inlined: a decode step frees one + tensor per request here, and the three-level call chain per uuid was + most of the cost. Same steps in the same order as the per-uuid form.""" + if not collectable: + return + del registered # no transfer engine to unregister from + store = self.tensor_store + store.mark_forgotten(collectable) + shard_dims = self.uuid_to_shard_dim + edge_names = self.uuid_to_edge_name + arena_ts = self._arena_ts + arena_locs = self._arena_locs + shm_files = self._shm_files + present = store.check_uuid_presence + remove = store.remove_tensor + for uuid in collectable: + try: + shard_dims.pop(uuid, None) + edge_names.pop(uuid, None) + if present(uuid): + remove(uuid) + else: + logger.warning( + "Trying to cleanup tensor %s, but uuid not found", uuid, + ) + arena_ts.pop(uuid, None) + if (loc := arena_locs.pop(uuid, None)) is not None: + self._arena.free(*loc) + if (path := shm_files.pop(uuid, None)) is not None: + try: + os.unlink(path) + except FileNotFoundError: + pass + except Exception as e: + logger.error( + "Error cleaning up tensor uuid %d: %s, skipping.", uuid, e, + ) + def _cleanup_by_uuid(self, uuid: int, registered: bool | None = None): # Grandparent cleanup (refcounts): skip the file manager's unlink. super(SharedMemoryCommunicationManager, self)._cleanup_by_uuid( From 50baf43e6a5dbd2c36ed17c946a244c25fa9a7ae Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:21:21 -0700 Subject: [PATCH 070/219] worker: send only the request entries that carry counts or stream consumption --- mstar/worker/worker.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 51bb3d05f9..733d942073 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -1626,6 +1626,23 @@ def _stream_consumption(self, rid: int) -> dict[str, int]: for edge_name, sbuf in req_info.stream_buffers.items() } + def _stream_consumption_batch( + self, rids: list[int], + ) -> list[tuple[int, dict[str, int]]]: + """``_stream_consumption`` for the rids that have stream buffers at + all; the rest would report an empty dict.""" + infos = self.request_state.per_request_info + out: list[tuple[int, dict[str, int]]] = [] + for rid in rids: + req_info = infos.get(rid) + if req_info is None or not req_info.stream_buffers: + continue + out.append((rid, { + edge_name: sbuf._consumed + for edge_name, sbuf in req_info.stream_buffers.items() + })) + return out + def _profiling_payloads(self, rids: list[int]) -> ParallelList: """rx/tx/timings for each rid's WORKER_GRAPHS_DONE.""" return ParallelList(rids, [ @@ -3129,6 +3146,13 @@ def _pp_stage(name: str) -> None: rid, batch_N.partition, published, ) buffered_publish.add(rid) + # Only the rids with something to report cross: a decode step has no + # new-token counts (the inline token is not counted) and no stream + # buffers, and an empty entry per rid cost a dict each to build and to + # extract on the runtime side. A rid left out gets nothing added, + # which is what an empty entry did. + ntc_rids = [rid for rid in send_rids if new_token_counts.get(rid)] + consumed = self._stream_consumption_batch(send_rids) send_input = SendInput( completion_id=route_output.completion_id, per_request_info=ParallelList( @@ -3139,12 +3163,10 @@ def _pp_stage(name: str) -> None: ], ), new_token_counts=ParallelList( - send_rids, - [new_token_counts.get(rid, {}) for rid in send_rids], + ntc_rids, [new_token_counts[rid] for rid in ntc_rids], ), stream_tokens_consumed=ParallelList( - send_rids, - [self._stream_consumption(rid) for rid in send_rids], + [rid for rid, _ in consumed], [c for _, c in consumed], ), profiling=self._profiling_payloads(send_rids) if self.enable_prof else None, From 54042512464aa1413b0da2bc39a78907c9ff9b60 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:21:21 -0700 Subject: [PATCH 071/219] engine: direct properties for the forward-pass wrapper's hot attributes --- mstar/engine/engine.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 450b1e1010..21ac348586 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -209,6 +209,29 @@ class ForwardPassInfoWrapper: info: CurrentForwardPassInfo metadata: InputMetadata + # The attributes every per-row step reads (prepare_inputs, the stop + # check) resolve without the ``__getattr__`` fallback, which costs a + # failed lookup plus a Python call per read. + @property + def dynamic_loop_iter_counts(self): + return self.metadata.dynamic_loop_iter_counts + + @property + def resource_configs(self): + return self.info.resource_configs + + @property + def max_tokens(self): + return self.info.max_tokens + + @property + def graph_walk(self): + return self.info.graph_walk + + @property + def step_metadata(self): + return self.info.step_metadata + def __getattr__(self, name): if name == "dynamic_loop_iter_counts": return self.metadata.dynamic_loop_iter_counts From bcfe673ab69fa51c31c9e6adaeb2bd9404c912bd Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:21:21 -0700 Subject: [PATCH 072/219] test: batched stream consumption report --- test/modular/test_send_prep_entries.py | 31 ++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 test/modular/test_send_prep_entries.py diff --git a/test/modular/test_send_prep_entries.py b/test/modular/test_send_prep_entries.py new file mode 100644 index 0000000000..22992ef3c5 --- /dev/null +++ b/test/modular/test_send_prep_entries.py @@ -0,0 +1,31 @@ +"""The send stage reports stream consumption only for requests that have +stream buffers; the rest would carry an empty dict each.""" +from types import SimpleNamespace + +from mstar.worker.worker import Worker + + +class _Buf: + def __init__(self, consumed): + self._consumed = consumed + + +def _worker(infos): + w = Worker.__new__(Worker) + w.request_state = SimpleNamespace(per_request_info=infos) + return w + + +def test_only_requests_with_stream_buffers_are_reported(): + infos = { + 1: SimpleNamespace(stream_buffers={}), + 2: SimpleNamespace(stream_buffers={"audio": _Buf(7), "text": _Buf(3)}), + 4: SimpleNamespace(stream_buffers={"audio": _Buf(0)}), + } + w = _worker(infos) + out = w._stream_consumption_batch([1, 2, 3, 4]) + assert out == [(2, {"audio": 7, "text": 3}), (4, {"audio": 0})] + # the per-request form agrees entry by entry + for rid, consumed in out: + assert w._stream_consumption(rid) == consumed + assert w._stream_consumption(1) == {} and w._stream_consumption(3) == {} From 75f5ea40e3033b2072cdff92e72078cfb3583821 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:21:21 -0700 Subject: [PATCH 073/219] ci: list the send prep test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 52e9bb4408..5f0a932622 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -53,3 +53,4 @@ test/modular/test_lazy_per_rid.py test/modular/test_step_context_rows.py test/modular/test_publish_snapshot_batch.py test/modular/test_sampler_for_cache.py +test/modular/test_send_prep_entries.py From c72d6af3fc32dd1adacc4dfea66e2324976edf58 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:25:58 -0700 Subject: [PATCH 074/219] kv commit: per-row constants hoisted, index offer only with a chain --- mstar/engine/resources/kv/manager.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 42e4d2188b..c3d82711dd 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -1320,6 +1320,10 @@ def clear_preplan(self): self._cached_plan_output = None def commit(self, step: KVStep, ctx: StepContext): + streams = self._streams + retention = step.retention + # the index only sees real rows of a non-capture step; checked once + index_on = self._index is not None and not ctx.capture # atomic against admit_retrieve reading stored_len on another thread with self._lock: # committed, so there is nothing left to unwind @@ -1327,7 +1331,7 @@ def commit(self, step: KVStep, ctx: StepContext): for segment in step.segments: if ctx.is_padding_row(segment.request_id): continue # ran against SINK_PAGE, holds nothing - stream = self._streams[segment.request_id][segment.label] + stream = streams[segment.request_id][segment.label] # cleared before the `step.commit` test: a step that keeps no # tokens (image_gen, action_gen) still read these pages, and # leaving the mark set would make the request unevictable @@ -1347,10 +1351,14 @@ def commit(self, step: KVStep, ctx: StepContext): stream.stored_len += segment.span # committed, so there is no refused admit left for a re-probe to answer stream.converted = False - policy = step.retention.get((segment.request_id, segment.label)) + policy = ( + retention.get((segment.request_id, segment.label)) + if retention else None + ) if policy is not None: self._adopt_retention(stream, policy, segment) - self._index_filled_pages(segment, stream, ctx) + if index_on and stream.chain is not None: + self._index_filled_pages(segment, stream, ctx) # so a claim taken in a window the mark misses still fails # `_commit_offload`'s generation guard stream.generation += 1 From 3a31d8b85f53bb0d6a432e7b9e6bdf736b235234 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:38:54 -0700 Subject: [PATCH 075/219] api_server: wake the streaming handlers once per message pass --- mstar/api_server/entrypoint.py | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 3b347bc40b..6c9e28e361 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -201,6 +201,32 @@ def notify(self) -> None: pass # the loop is closed: the client is gone +def _set_events(wakes: list) -> None: + for wake in wakes: + wake.set() + + +def notify_many(reqs) -> None: + """``PendingRequest.notify`` for a batch: one loop wake-up (a self-pipe + write) per event loop rather than one per chunk. A step's RESULT_TOKENS + frame lands one chunk per running request, so the message thread was + waking the loop once per token.""" + by_loop: dict[int, tuple[Any, list]] = {} + for req in reqs: + loop, wake = req.loop, req.wake + if loop is None or wake is None: + continue + entry = by_loop.get(id(loop)) + if entry is None: + entry = by_loop[id(loop)] = (loop, []) + entry[1].append(wake) + for loop, wakes in by_loop.values(): + try: + loop.call_soon_threadsafe(_set_events, wakes) + except RuntimeError: + pass # the loop is closed: the client is gone + + def _chunk_to_ndjson_payload(chunk: ResultChunk) -> str: """Serialize one result chunk as an NDJSON line.""" return json.dumps({ @@ -653,6 +679,9 @@ def _process_messages(self) -> None: update.preprocess_finish_time req.profile.inputs = update.inputs + # the handlers are woken once per pass, after every chunk of + # the pass has been appended, not once per chunk + to_wake: dict[int, PendingRequest] = {} for result_chunk in self.preprocess_worker.get_result_chunks(): logger.debug( "Got result chunk of %s modality for request %s", @@ -674,7 +703,7 @@ def _process_messages(self) -> None: req.profile.timing.first_chunk_time = now req.profile.timing.last_chunk_time = now req.chunks.append(result_chunk) - req.notify() + to_wake[id(req)] = req if result_chunk.modality == "error": # The data worker failed this request (preprocess, @@ -687,11 +716,12 @@ def _process_messages(self) -> None: (result_chunk.metadata or {}).get("status", 500) ) req.event.set() - req.notify() # Park the rid so _prune_recently_completed reclaims # the data worker's per-request state once the # client lets go of the request. self.recently_completed[rid] = time.time() + if to_wake: + notify_many(to_wake.values()) except Exception: if self.running: logger.exception("Error in message processing loop") From b861dcfd15db9622cb1ca5487ee98906d88a0156 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 18:38:54 -0700 Subject: [PATCH 076/219] test: batched handler wake-up --- test/modular/test_api_result_delivery.py | 29 ++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/test/modular/test_api_result_delivery.py b/test/modular/test_api_result_delivery.py index 92901871f5..d90109f3fa 100644 --- a/test/modular/test_api_result_delivery.py +++ b/test/modular/test_api_result_delivery.py @@ -274,3 +274,32 @@ def producer(): # each chunk reached the handler within the 50 ms backstop, not a poll later for (_, t_got), t_put in zip(got, times, strict=True): assert t_got - t_put < 0.05, f"chunk took {1000 * (t_got - t_put):.1f} ms" + + +def test_notify_many_wakes_every_handler_with_one_loop_call(): + """A pass that delivered chunks to many requests wakes their handlers + through one call into the loop, and skips requests with no handler.""" + from mstar.api_server.entrypoint import notify_many + + class _Loop: + def __init__(self): + self.calls = [] + + def call_soon_threadsafe(self, fn, *args): + self.calls.append((fn, args)) + fn(*args) + + class _Closed: + def call_soon_threadsafe(self, fn, *args): + raise RuntimeError("Event loop is closed") + + loop = _Loop() + reqs = [_pending_request() for _ in range(4)] + for req in reqs[:3]: + req.loop, req.wake = loop, asyncio.Event() + reqs[3].loop, reqs[3].wake = _Closed(), asyncio.Event() + orphan = _pending_request() # no handler attached yet + notify_many(reqs + [orphan]) + assert len(loop.calls) == 1 + assert all(req.wake.is_set() for req in reqs[:3]) + assert not reqs[3].wake.is_set() From f5f5303b3c9acb67219661ff6545ed2e236b4040 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:01:49 -0700 Subject: [PATCH 077/219] rust runtime: env-gated timing of send_outputs (body vs GIL reacquire) --- rust/src/graph/runtime.rs | 56 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index 1af202fc52..e8748f1466 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -2928,7 +2928,10 @@ impl GraphRuntime { // Python thread may call into this runtime while a send is in // flight: it would get `Already mutably borrowed` rather than block. // Only the worker's main loop does today. - py.allow_threads(move || -> PyResult> { + let t_call = std::time::Instant::now(); + let out = py.allow_threads(move || -> PyResult> { + let t_body = std::time::Instant::now(); + let out = (move || -> PyResult> { let request_infos: Vec<(u32, Option>)> = info_rids.into_iter().zip(request_infos).collect(); let new_token_counts: Vec<(u32, Vec<(String, i64)>)> = ntc_rids @@ -3110,7 +3113,12 @@ impl GraphRuntime { sent_rids.push(rid); } Ok(sent_rids) - }) + })(); + send_timing::record_body(t_body); + out + }); + send_timing::record_total(t_call); + out } fn get_loop_stop_times( @@ -3434,3 +3442,47 @@ mod tests { assert!(!rids.release(42)); } } + +/// Env-gated timing of `send_outputs` (MSTAR_RUST_SEND_TIMING=1): how long +/// the GIL-free body takes against the whole call, which also pays for +/// taking the GIL back. Printed to stderr every 500 calls. +mod send_timing { + use std::sync::OnceLock; + use std::sync::atomic::{AtomicU64, Ordering}; + + static ENABLED: OnceLock = OnceLock::new(); + static CALLS: AtomicU64 = AtomicU64::new(0); + static BODY_NS: AtomicU64 = AtomicU64::new(0); + static TOTAL_NS: AtomicU64 = AtomicU64::new(0); + const EVERY: u64 = 500; + + fn enabled() -> bool { + *ENABLED.get_or_init(|| { + std::env::var("MSTAR_RUST_SEND_TIMING").map(|v| v == "1").unwrap_or(false) + }) + } + + pub fn record_body(t: std::time::Instant) { + if enabled() { + BODY_NS.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed); + } + } + + pub fn record_total(t: std::time::Instant) { + if !enabled() { + return; + } + TOTAL_NS.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed); + let n = CALLS.fetch_add(1, Ordering::Relaxed) + 1; + if n % EVERY == 0 { + let body = BODY_NS.swap(0, Ordering::Relaxed) as f64; + let total = TOTAL_NS.swap(0, Ordering::Relaxed) as f64; + let per = EVERY as f64 * 1000.0; + eprintln!( + "mstar_rust send_outputs timing over {} calls: body {:.1} us, \ + total {:.1} us, gil reacquire {:.1} us", + EVERY, body / per, total / per, (total - body) / per, + ); + } + } +} From 6de442383a429e9ee0387a035dc868d6bb04b58a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:09:51 -0700 Subject: [PATCH 078/219] worker: MSTAR_SPEC_YIELD_HOLD_MS holds the decode loop's yield so arrivals prefill together --- mstar/worker/worker.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 733d942073..70b58157ab 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -3694,6 +3694,11 @@ def run(self) -> None: spec_peek_for_fairness = ( os.environ.get("MSTAR_SPEC_PEEK_FOR_FAIRNESS", "1") == "1" ) + # How long the loop keeps speculating once another step is ready, so + # that at high concurrency the requests arriving over that window + # prefill in one step rather than one step each. 0 yields at once. + spec_yield_hold_s = float(os.environ.get("MSTAR_SPEC_YIELD_HOLD_MS", "0")) / 1000.0 + fairness_ready_since: list[float | None] = [None] consecutive_spec_steps = 0 yield_away_from_target: tuple[str, str] | None = None # Build N+2 while N+1 runs, between the two halves of N's @@ -3725,7 +3730,7 @@ def _build_speculation( # another (node, walk) actually ready to schedule on # this worker. On single-walk workers (Orpheus LLM, # Orpheus SNAC) this returns False and we always speculate. - must_yield_for_fairness = ( + other_ready = ( spec_peek_for_fairness and consecutive >= 1 and self.scheduler.has_ready_excluding( @@ -3733,10 +3738,22 @@ def _build_speculation( (pending.node_name, pending.graph_walk), ) ) + if other_ready: + now = _time.perf_counter() + if fairness_ready_since[0] is None: + fairness_ready_since[0] = now + must_yield_for_fairness = ( + now - fairness_ready_since[0] >= spec_yield_hold_s + ) + else: + fairness_ready_since[0] = None + must_yield_for_fairness = False must_yield_away = ( consecutive >= max_consecutive_spec or must_yield_for_fairness ) + if must_yield_away: + fairness_ready_since[0] = None if not must_yield_away and not admit_refused: if self.enable_nvtx: range_push("worker.speculate", synchronize=False) From 2b422b903d087e1c5d30137a34802a8331ae75d3 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:09:51 -0700 Subject: [PATCH 079/219] docs: MSTAR_SPEC_YIELD_HOLD_MS --- docs/environment_variables.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 5a50265888..e1bad4b84b 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -323,6 +323,15 @@ Worker scheduling - ``1024`` - Cap on back-to-back speculative steps before the leader yields to other ready work. + * - ``MSTAR_SPEC_YIELD_HOLD_MS`` + - ``0`` + - How long (ms) the speculative decode loop keeps going after another + step becomes ready, typically the prefill of requests that just + arrived. With ``0`` the loop yields at the next iteration and every + arrival tends to get a prefill step of its own; a hold of a few + milliseconds lets the arrivals of that window prefill together, which + is worth decode throughput at high concurrency and costs up to that + much time to first token. * - ``MSTAR_SPEC_PEEK_FOR_FAIRNESS`` - ``1`` - Yield the speculation chain only when another (node, walk) is From 720587d4dbbe37c617384e59deb12d9d0683db78 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:13:29 -0700 Subject: [PATCH 080/219] rust runtime: flat export of the dynamic loop iterations --- rust/src/graph/runtime.rs | 43 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index e8748f1466..31da54d74a 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -1866,6 +1866,49 @@ impl GraphRuntime { .collect() } + /// `get_dynamic_loop_iters` as four flat columns: the distinct loop names + /// once, then per entry the request's index in `request_ids`, the name's + /// index and the iteration. One string per loop name instead of one per + /// request per step. + fn get_dynamic_loop_iters_flat( + &self, request_ids: Vec, partition: &str, + ) -> (Vec, Vec, Vec, Vec) { + let mut names: Vec = Vec::new(); + let mut name_idx_of: FxHashMap = FxHashMap::default(); + let mut rid_idx: Vec = Vec::new(); + let mut name_idx: Vec = Vec::new(); + let mut iters: Vec = Vec::new(); + let Some(p) = self.interner.get(partition) else { + return (names, rid_idx, name_idx, iters); + }; + for (i, &rid) in request_ids.iter().enumerate() { + let Some(info) = self.info(rid) else { continue }; + let Some(part) = info.partitions.get(&p) else { continue }; + for &wg in &part.walk_worker_graphs { + let Some(state) = self.state(wg, rid) else { + continue; + }; + let g = &self.graphs[wg as usize]; + for (lid, iter) in state.loop_indices().into_iter().enumerate() { + let sym = g.lp(lid as LoopId).name; + let ni = match name_idx_of.get(&sym) { + Some(&n) => n, + None => { + let n = names.len() as u32; + names.push(self.interner.name(sym).to_string()); + name_idx_of.insert(sym, n); + n + } + }; + rid_idx.push(i as u32); + name_idx.push(ni); + iters.push(iter); + } + } + } + (names, rid_idx, name_idx, iters) + } + /// No-op by construction. Python clears `tensor_info` off the node's /// output edges because the edges are per-request objects that carry it; /// here outputs are passed into `complete_and_route_batch`, so there is From d47c22fe239616c3b3763efeab4f1dc613241e4d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:13:29 -0700 Subject: [PATCH 081/219] graph runtime: update_dynamic_loop_iters hook --- mstar/graph/runtime/base.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/mstar/graph/runtime/base.py b/mstar/graph/runtime/base.py index 2f77e613d2..c89786ddc9 100644 --- a/mstar/graph/runtime/base.py +++ b/mstar/graph/runtime/base.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from dataclasses import dataclass, field from typing import NamedTuple @@ -598,6 +599,18 @@ def get_dynamic_loop_iters( ) -> ParallelList[int, dict[str,int]]: pass + def update_dynamic_loop_iters( + self, per_request_info: Mapping[int, "CurrentForwardPassInfo"], partition: str, + ) -> None: + """Refresh ``dynamic_loop_iter_counts`` on each request's forward-pass + info from the runtime's loop counters. The default goes through + ``get_dynamic_loop_iters``; a runtime with a cheaper bulk form + overrides this.""" + for rid, new_iters in self.get_dynamic_loop_iters( + list(per_request_info), partition=partition, + ): + per_request_info[rid].dynamic_loop_iter_counts.update(new_iters) + @abstractmethod def is_async_schedulable(self, node_name: str, graph_walk: str) -> bool: """Whether this node opts into async scheduling. Structural, so it From 880cdbe03395bc2b4e5ab4da2b80de53b00f5beb Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:13:29 -0700 Subject: [PATCH 082/219] rust runtime wrapper: refresh loop iterations from the flat export --- mstar/graph/runtime/rust.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mstar/graph/runtime/rust.py b/mstar/graph/runtime/rust.py index 386d7dbc9c..b9edc253a2 100644 --- a/mstar/graph/runtime/rust.py +++ b/mstar/graph/runtime/rust.py @@ -395,6 +395,18 @@ def ingest_inputs_batch( signals, can_buffer, is_streaming, ) + def update_dynamic_loop_iters( + self, per_request_info, partition: str, + ) -> None: + rids = list(per_request_info) + names, rid_idx, name_idx, iters = self._rust.get_dynamic_loop_iters_flat( + rids, partition, + ) + for k in range(len(iters)): + per_request_info[rids[rid_idx[k]]].dynamic_loop_iter_counts[ + names[name_idx[k]] + ] = iters[k] + def get_dynamic_loop_iters( self, request_ids: list[int], partition: str, ) -> ParallelList[int, dict[str, int]]: From cb0b9cef5a14ecedcea652d1fee4c5375b923f14 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:13:30 -0700 Subject: [PATCH 083/219] worker: refresh loop iterations through the runtime hook --- mstar/worker/worker.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 70b58157ab..c529fdb2d6 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -3032,11 +3032,9 @@ def _pp_stage(name: str) -> None: "worker.postprocess.route", _time.perf_counter() - _t_route, ) - per_request_info = batch_N.node_batch.per_request_info - for rid, new_iters in self._graph_runtime.get_dynamic_loop_iters( - list(per_request_info), partition=batch_N.partition, - ): - per_request_info[rid].dynamic_loop_iter_counts.update(new_iters) + self._graph_runtime.update_dynamic_loop_iters( + batch_N.node_batch.per_request_info, batch_N.partition, + ) # Normally empty: cleanup_consumed_inputs ran above and took them. Not # empty if a completion ever precedes it, and then nobody else will. From aa5c467e0e94fe09663a804fd612571a9b60d838 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:13:30 -0700 Subject: [PATCH 084/219] test: update_dynamic_loop_iters matches the per-rid form --- test/rust/test_rust_graph_runtime.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/test/rust/test_rust_graph_runtime.py b/test/rust/test_rust_graph_runtime.py index f5b274006d..1faec35751 100644 --- a/test/rust/test_rust_graph_runtime.py +++ b/test/rust/test_rust_graph_runtime.py @@ -459,6 +459,20 @@ def test_dynamic_loop_iters_for_an_unknown_partition_are_empty(runtime): assert runtime.get_dynamic_loop_iters([rid], "nope").values == [{}] +def test_update_dynamic_loop_iters_matches_the_per_rid_form(runtime): + from types import SimpleNamespace + rid = _admit(runtime) + infos = {rid: SimpleNamespace(dynamic_loop_iter_counts={"stale": 7})} + runtime.update_dynamic_loop_iters(infos, "default") + expect = dict(runtime.get_dynamic_loop_iters([rid], "default").values[0]) + expect["stale"] = 7 + assert infos[rid].dynamic_loop_iter_counts == expect + infos = {rid: SimpleNamespace(dynamic_loop_iter_counts={})} + runtime.update_dynamic_loop_iters(infos, "nope") + assert infos[rid].dynamic_loop_iter_counts == {} + runtime.update_dynamic_loop_iters({}, "default") + + # --- scheduling -------------------------------------------------------------- def test_a_ready_node_is_reported(runtime): From 92bed840bafb17eb45264d07d4b06c02bc6a89b6 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:23 -0700 Subject: [PATCH 085/219] kv: batched prefix-chain extension off the stop check's host rows --- mstar/engine/resources/kv/manager.py | 37 ++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index c3d82711dd..1204b7513d 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -706,6 +706,43 @@ def extend_prefix_chain( return stream.chain.extend(sampled[0].flatten().tolist(), self.config.page_size) + def extend_prefix_chains_batch( + self, request_ids: list[str], node_name: str, graph_walk: str, + host_rows, + ) -> None: + """``extend_prefix_chain`` for a step, from the stop check's row + buffers (``HostRows``: row i of every buffer belongs to + ``host_rows.request_ids[i]``). One lock, one ``tolist`` per buffer + instead of one per request.""" + row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + buffers = host_rows.buffers + rows_of: dict[str, list] = {} + page_size = self.config.page_size + with self._lock: + for rid in request_ids: + label = self._keyed_label(rid, node_name, graph_walk) + if label is None: + continue + tensor = (self._overrides[rid].prefix_decode or {}).get(label) + if not tensor: + continue + rows = rows_of.get(tensor) + if rows is None: + buf = buffers.get(tensor) + if not torch.is_tensor(buf) or buf.dim() == 0: + continue + rows = rows_of[tensor] = buf.reshape(buf.shape[0], -1).tolist() + i = row_of.get(rid) + if i is None or i >= len(rows): + continue + stream = self._streams.get(rid, {}).get(label) + if stream is None or stream.chain is None or stream.chain.unkeyed is None: + continue + if stream.released: + stream.chain = None + continue + stream.chain.extend(rows[i], page_size) + def _release_lease(self, stream: CacheStream) -> None: """Give back a lease `admit` never converted; a converted one is released with `page_indices`.""" From cb338d8352d09031f23983f474e7e18ff7fd05fd Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:23 -0700 Subject: [PATCH 086/219] runner: batched chain extension; sweep resources that override the hook without the opt-in --- mstar/engine/resources/runner.py | 36 ++++++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index 14d54435e9..dbc1670c46 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -139,10 +139,21 @@ def prefix_chain_keys(self, node_name: str | None) -> list[str]: """The node's resources with a prefix cache open right now. Empty in the common case, which lets the caller skip a per-request sweep that was six no-op calls per request per step.""" - return [ - key for key in self._sweep(self._node_order, self._order, node_name) - if self._resources[key].keys_prefix_chains - ] + keys: list[str] = [] + for key in self._sweep(self._node_order, self._order, node_name): + resource = self._resources[key] + if resource.keys_prefix_chains: + keys.append(key) + continue + # a resource that overrides the hook without saying when it keys + # is swept as before the skip existed + cls = type(resource) + if ( + cls.extend_prefix_chain is not Resource.extend_prefix_chain + and cls.keys_prefix_chains is Resource.keys_prefix_chains + ): + keys.append(key) + return keys def extend_prefix_chains( self, rid: str, node_name: str, graph_walk: str, outputs, @@ -156,6 +167,23 @@ def extend_prefix_chains( rid, node_name, graph_walk, outputs, ) + def extend_prefix_chains_batch( + self, request_ids: list[str], node_name: str, graph_walk: str, + host_rows, outputs, keys: list[str], + ) -> None: + """``extend_prefix_chains`` for a step: a resource with a batch form + takes the host rows; the others are offered each request's host copy.""" + for key in keys: + resource = self._resources[key] + batched = getattr(resource, "extend_prefix_chains_batch", None) + if batched is not None: + batched(request_ids, node_name, graph_walk, host_rows) + continue + for rid in request_ids: + per_rid = outputs.get(rid) + if isinstance(per_rid, dict): + resource.extend_prefix_chain(rid, node_name, graph_walk, per_rid) + def _check_preplan_deps(self) -> None: """A pre-planning resource's dependencies must pre-plan too. From da6df5b78624de242b788c1fd71c98596d36465d Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:24 -0700 Subject: [PATCH 087/219] resources: document when the chain hook is swept --- mstar/engine/resources/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/base.py b/mstar/engine/resources/base.py index c304ad4363..a72dddb770 100644 --- a/mstar/engine/resources/base.py +++ b/mstar/engine/resources/base.py @@ -131,7 +131,9 @@ def extend_prefix_chain( @property def keys_prefix_chains(self) -> bool: """Whether ``extend_prefix_chain`` can do anything right now (a prefix - cache is open). The runner skips the per-request sweep otherwise.""" + cache is open). The runner skips the per-request sweep otherwise. A + resource that overrides ``extend_prefix_chain`` but not this is swept + every step.""" return False def fingerprint(self) -> bytes | None: From 61ea439e141974077b3b52e21ae1be9e03eb0c6c Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:24 -0700 Subject: [PATCH 088/219] engine: extend prefix chains from the host rows --- mstar/engine/engine.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 21ac348586..d3995e966d 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -796,12 +796,19 @@ def inline_client_signals( def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], + host_rows: HostRows | None = None, ) -> None: """Key what this step generated, from the stop check's host copy.""" keys = self._runner.prefix_chain_keys(batch.node_name) if not keys: return walk = batch.step_context.graph_walk + if host_rows is not None: + self._runner.extend_prefix_chains_batch( + batch.request_ids, batch.node_name, walk, host_rows, outputs, + keys=keys, + ) + return for rid in batch.request_ids: per_rid = outputs.get(rid) if isinstance(per_rid, dict): From f66cb9b06c942fa11655a0f2808ee2060aa1a42f Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:24 -0700 Subject: [PATCH 089/219] worker: hand the host rows to the prefix-chain extension --- mstar/worker/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index c529fdb2d6..ce8154cae2 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2878,7 +2878,7 @@ def _pp_stage(name: str) -> None: "worker.postprocess.check_stop", _time.perf_counter() - _t_stop, ) # the same host copy, before stops, so a request ending here still indexes its pages - engine.extend_prefix_chains(batch_N.node_batch, cpu_outputs) + engine.extend_prefix_chains(batch_N.node_batch, cpu_outputs, host_rows=host_rows) # Stream-terminated loop: a stream-consuming node inside a loop has no # internal stop signal (unlike a self-EOS loop), so when it consumes the From a490f5b7abd2e58675eeaedfe057e6a8f81a0aee Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:24 -0700 Subject: [PATCH 090/219] test: batched prefix-chain extension --- test/modular/test_kv_prefix_batch.py | 109 +++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 test/modular/test_kv_prefix_batch.py diff --git a/test/modular/test_kv_prefix_batch.py b/test/modular/test_kv_prefix_batch.py new file mode 100644 index 0000000000..c0f413008e --- /dev/null +++ b/test/modular/test_kv_prefix_batch.py @@ -0,0 +1,109 @@ +"""The batched prefix-chain extension (one call per step off the stop check's +host rows) keys exactly what the per-request form keys.""" +from __future__ import annotations + +import importlib.util +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, ".") + +from mstar.engine.resources.runner import StepRunner # noqa: E402 +from mstar.model.submodule_base import HostRows # noqa: E402 + +_spec = importlib.util.spec_from_file_location( + "kv_prefix_decode_helpers", pathlib.Path(__file__).with_name("test_kv_prefix_decode.py"), +) +H = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(H) + + +@pytest.fixture(autouse=True) +def _stub_transfer(monkeypatch): + monkeypatch.setattr(H.manager_mod, "KVTransferManager", H._StubTransfer) + + +def _chain_state(kv, rid): + stream = kv._streams[rid]["main"] + chain = stream.chain + if chain is None: + return None + return (chain.covered_len, list(chain.keys) if hasattr(chain, "keys") else None, + chain.unkeyed if hasattr(chain, "unkeyed") else None) + + +def _run(batched: bool): + kv = H._manager() + rids = ["r0", "r1"] + prompts = {"r0": list(range(100)), "r1": list(range(50, 50 + H.PAGE_SIZE + 10))} + for rid in rids: + H._ingest(kv, rid, prompts[rid]) + # the prefill of each, and the token it sampled + for rid in rids: + H._step(kv, rid, len(prompts[rid])) + sampled = {"r0": 9000, "r1": 9100} + if batched: + rows = HostRows(request_ids=tuple(rids), buffers={ + H.TENSOR: torch.tensor([[sampled[r]] for r in rids]), + }) + kv.extend_prefix_chains_batch(rids, H.NODE, H.WALK, rows) + else: + for rid in rids: + kv.extend_prefix_chain(rid, H.NODE, H.WALK, {H.TENSOR: [torch.tensor([sampled[rid]])]}) + # decode steps: each writes the id sampled before it and samples a new one + for t in range(1, 2 * H.PAGE_SIZE + 3): + for rid in rids: + H._step(kv, rid, 1) + tokens = {"r0": 9000 + t, "r1": 9100 + t} + if batched: + rows = HostRows(request_ids=tuple(reversed(rids)), buffers={ + H.TENSOR: torch.tensor([[tokens[r]] for r in reversed(rids)]), + }) + kv.extend_prefix_chains_batch(rids, H.NODE, H.WALK, rows) + else: + for rid in rids: + kv.extend_prefix_chain(rid, H.NODE, H.WALK, {H.TENSOR: [torch.tensor([tokens[rid]])]}) + kv.assert_pages_conserved() + return H._indexed(kv), {rid: _chain_state(kv, rid) for rid in rids} + + +def test_batched_extension_keys_the_same_pages(): + per_rid = _run(batched=False) + batched = _run(batched=True) + assert per_rid == batched + assert per_rid[0] >= 3, "the generated pages were not indexed" + + +def test_batched_extension_skips_rows_it_does_not_have(): + kv = H._manager() + prompt = list(range(H.PAGE_SIZE)) + H._ingest(kv, "r0", prompt) + H._step(kv, "r0", len(prompt)) + before = _chain_state(kv, "r0") + # no buffer under the keyed tensor name, and a request with no row + kv.extend_prefix_chains_batch(["r0", "ghost"], H.NODE, H.WALK, + HostRows(request_ids=("r0",), buffers={"other": torch.tensor([[1]])})) + kv.extend_prefix_chains_batch(["r0"], H.NODE, H.WALK, + HostRows(request_ids=("r1",), buffers={H.TENSOR: torch.tensor([[1]])})) + assert _chain_state(kv, "r0") == before + + +def test_runner_dispatch_falls_back_per_request(): + calls = [] + + class _Batch: + def extend_prefix_chains_batch(self, rids, node, walk, host_rows): + calls.append(("batch", tuple(rids))) + + class _PerRid: + def extend_prefix_chain(self, rid, node, walk, outputs): + calls.append(("rid", rid, outputs)) + + runner = StepRunner.__new__(StepRunner) + runner._resources = {"a": _Batch(), "b": _PerRid()} + outputs = {"r0": {"x": 1}, "r1": None} + runner.extend_prefix_chains_batch(["r0", "r1"], "n", "w", HostRows((), {}), outputs, keys=["a", "b"]) + assert calls == [("batch", ("r0", "r1")), ("rid", "r0", {"x": 1})] From b369e75b1bea1c273525e8cafd7d29d5f18e4a68 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:18:24 -0700 Subject: [PATCH 091/219] ci: list the prefix batch test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 5f0a932622..94984fced4 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -54,3 +54,4 @@ test/modular/test_step_context_rows.py test/modular/test_publish_snapshot_batch.py test/modular/test_sampler_for_cache.py test/modular/test_send_prep_entries.py +test/modular/test_kv_prefix_batch.py From 60c0ab27d3bf136eb0301e42252f505406d89fcc Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 19:30:57 -0700 Subject: [PATCH 092/219] engine: read the keyed walks leniently in prepare_inputs --- mstar/engine/engine.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index d3995e966d..4e54522116 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -750,8 +750,10 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: tensors = batch.per_request_input_tensors metadata = batch.per_request_input_metadata final_rids = batch.final_stream_rids - # only a keyed walk probes the prefix cache (see _skip_cached_prefix) - probe_prefix = walk in self._keyed_walks.get(batch.node_name, ()) + # only a keyed walk probes the prefix cache (see _skip_cached_prefix); + # read leniently, as the probe itself was only reached per prepared row + keyed = getattr(self, "_keyed_walks", None) + probe_prefix = bool(keyed) and walk in keyed.get(batch.node_name, ()) for rid in batch.request_ids: try: req_inputs = prepare( From d35c66b272dd4fd33ae1c174140a33c217a1f83a Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 20:02:00 -0700 Subject: [PATCH 093/219] zmq communicator: non-blocking try_send --- rust/src/communicator.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/rust/src/communicator.rs b/rust/src/communicator.rs index c0b4a7c1d1..c15b7de7d3 100644 --- a/rust/src/communicator.rs +++ b/rust/src/communicator.rs @@ -233,6 +233,36 @@ impl RawZmqCommunicator { Ok(()) } + /// `send` that never blocks: Ok(true) when the frame was queued, Ok(false) + /// when the peer is at its high-water mark (the caller then falls back to + /// the blocking `send`, with the GIL released if it holds it). + pub fn try_send(&self, peer_id: &str, payload: &[u8]) -> Result { + let socket = { + let mut peers = self.peers.lock().expect("peers lock"); + match peers.get(peer_id) { + Some(s) => s.clone(), + None => { + let endpoint = self.resolve(peer_id)?; + let push = self.ctx.socket(zmq::PUSH)?; + push.set_linger(0)?; + push.connect(&endpoint)?; + let s = std::sync::Arc::new(Mutex::new(push)); + peers.insert(peer_id.to_string(), s.clone()); + s + } + } + }; + let result = socket + .lock() + .expect("peer socket lock") + .send(payload, zmq::DONTWAIT); + match result { + Ok(()) => Ok(true), + Err(zmq::Error::EAGAIN) => Ok(false), + Err(e) => Err(e.into()), + } + } + /// Non-blocking: next inbound frame, or None. pub fn try_recv(&self) -> Option { let pull = self.pull.lock().expect("pull lock"); From 22278c5984d2915a051ac63a820ba07a290d21b7 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 20:02:00 -0700 Subject: [PATCH 094/219] rust runtime: send frames with the GIL held, release it only for a peer at its high-water mark --- rust/src/graph/runtime.rs | 83 ++++++++++++++++++++++++++++----------- 1 file changed, 61 insertions(+), 22 deletions(-) diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index 31da54d74a..c8c571fac7 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -273,12 +273,44 @@ pub struct GraphRuntime { /// is right for the no-transport case: the frame builders are exercised /// without a socket, and the worker's flag gate already refuses to pair /// the Rust runtime with a communicator it cannot share. - fn dispatch(&self, peer: &str, bytes: &[u8]) -> PyResult<()> { - if let Some(comm) = &self.communicator { - comm.send(peer, bytes) - .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + /// Send frames in order. Each is first tried without blocking, with the + /// GIL held: a send that goes through costs no GIL reacquisition, which + /// the released-GIL form paid at up to a millisecond per step on the + /// worker's main thread. From the first frame a peer cannot take (its + /// high-water mark), the rest go out blocking with the GIL released, as + /// every send did before, so a stalled peer still cannot freeze the + /// other Python threads. + fn send_frames( + &self, py: Python<'_>, frames: Vec<(String, Vec)>, + ) -> PyResult<()> { + let Some(comm) = &self.communicator else { return Ok(()) }; + let mut rest: Vec<(String, Vec)> = Vec::new(); + let mut blocked = false; + for (peer, bytes) in frames { + if blocked { + rest.push((peer, bytes)); + continue; + } + match comm.try_send(&peer, &bytes) { + Ok(true) => {} + Ok(false) => { + blocked = true; + rest.push((peer, bytes)); + } + Err(e) => return Err(PyRuntimeError::new_err(e.to_string())), + } } - Ok(()) + if rest.is_empty() { + return Ok(()); + } + let comm = comm.clone(); + py.allow_threads(move || -> Result<(), crate::communicator::CommError> { + for (peer, bytes) in &rest { + comm.send(peer, bytes)?; + } + Ok(()) + }) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) } /// One WORKER_GRAPHS_DONE frame. Drains the request's pending persist @@ -2810,7 +2842,9 @@ impl GraphRuntime { loop_names: Vec>, ) -> PyResult> { // One release for the whole batch; see send_outputs. - py.allow_threads(move || -> PyResult> { + // bookkeeping with the GIL held; the frames go out through send_frames + let mut frames: Vec<(String, Vec)> = Vec::new(); + let out = (|| -> PyResult> { if rids.len() != loop_names.len() { return Err(PyValueError::new_err( "stop_loops_batched: rids and loop_names must be the same length", @@ -2859,12 +2893,15 @@ impl GraphRuntime { loop_stop_times: stop_times.clone(), } .encode(&self.interner); - self.dispatch(self.interner.name(worker), &bytes)?; + frames.push((self.interner.name(worker).to_string(), bytes)); } stopped_rids.push(rid); } Ok(stopped_rids) - }) + })(); + let out = out?; + self.send_frames(py, frames)?; + Ok(out) } /// A peer's STOP_LOOPS landing here. @@ -2972,9 +3009,12 @@ impl GraphRuntime { // flight: it would get `Already mutably borrowed` rather than block. // Only the worker's main loop does today. let t_call = std::time::Instant::now(); - let out = py.allow_threads(move || -> PyResult> { - let t_body = std::time::Instant::now(); - let out = (move || -> PyResult> { + // Encoding and bookkeeping run with the GIL held (tens of + // microseconds); the frames go out through send_frames, which + // releases the GIL only for a peer at its high-water mark. + let mut frames: Vec<(String, Vec)> = Vec::new(); + let t_body = std::time::Instant::now(); + let out = (|| -> PyResult> { let request_infos: Vec<(u32, Option>)> = info_rids.into_iter().zip(request_infos).collect(); let new_token_counts: Vec<(u32, Vec<(String, i64)>)> = ntc_rids @@ -3059,9 +3099,8 @@ impl GraphRuntime { )); } } - for (worker, bytes) in &frames_out { - let name = self.interner.name(*worker).to_string(); - self.dispatch(&name, bytes)?; + for (worker, bytes) in frames_out { + frames.push((self.interner.name(worker).to_string(), bytes)); } for (rid, signal, modality, refs) in plan.emit { @@ -3086,7 +3125,7 @@ impl GraphRuntime { } .encode(&self.interner, bk.strings()) }; - self.dispatch("api_server", &bytes)?; + frames.push(("api_server".to_string(), bytes)); } // Inline emits: one RESULT_TOKENS frame per (signal, modality) @@ -3128,7 +3167,7 @@ impl GraphRuntime { modality: &modality, } .encode(&self.interner); - self.dispatch("api_server", &bytes)?; + frames.push(("api_server".to_string(), bytes)); } } @@ -3152,16 +3191,16 @@ impl GraphRuntime { profiling.get(&rid).and_then(|b| b.as_deref()), spec_flags.get(&rid).copied().unwrap_or(false), )?; - self.dispatch("conductor", &bytes)?; + frames.push(("conductor".to_string(), bytes)); sent_rids.push(rid); } Ok(sent_rids) - })(); - send_timing::record_body(t_body); - out - }); + })(); + send_timing::record_body(t_body); + let out = out?; + self.send_frames(py, frames)?; send_timing::record_total(t_call); - out + Ok(out) } fn get_loop_stop_times( From 04a8f369f331cd7830557ece34f955c45256a52f Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 20:08:15 -0700 Subject: [PATCH 095/219] docs: MSTAR_RUST_SEND_TIMING --- docs/environment_variables.rst | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index e1bad4b84b..e3a821e4fc 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -332,6 +332,12 @@ Worker scheduling milliseconds lets the arrivals of that window prefill together, which is worth decode throughput at high concurrency and costs up to that much time to first token. + * - ``MSTAR_RUST_SEND_TIMING`` + - unset + - ``1`` makes the Rust runtime print, every 500 calls, how long the + worker's per-step send spends in its own work against the whole call + (which includes taking the GIL back when the send had to release it). + A diagnostic for the host-overhead work; off by default. * - ``MSTAR_SPEC_PEEK_FOR_FAIRNESS`` - ``1`` - Yield the speculation chain only when another (node, walk) is From 5db372f206a429e9380e8f329ca42313293793e3 Mon Sep 17 00:00:00 2001 From: merceod Date: Wed, 7 Oct 2026 22:35:07 -0700 Subject: [PATCH 096/219] data worker: sort the import block --- mstar/api_server/data_worker.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index f34c75a5b2..152f2e121e 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -24,7 +24,6 @@ ResultChunk, ResultTensors, ) - from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig From 59127087676d511eafebdf10e7bb1f5db7adb3ba Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 04:37:42 -0700 Subject: [PATCH 097/219] rust runtime: the inline signal keys of a route input are optional --- rust/src/graph/runtime.rs | 52 +++++++++++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 5 deletions(-) diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index c8c571fac7..5a2e9756f8 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -11,7 +11,7 @@ use crate::tensors::{SharedBookkeeping, TensorBookkeeping}; use crate::communicator::RawZmqCommunicator; use crate::PyZmqCommunicator; use std::sync::Arc; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; use rustc_hash::{FxHashMap, FxHashSet}; @@ -1124,9 +1124,9 @@ pub struct LoopStopArg { #[pyo3(item)] wg_fwd_pass_idx: u32, } -/// `RouteInput`. +/// `RouteInput`, the keys every caller sends. #[derive(FromPyObject)] -pub struct RouteArg { +struct RouteArgRequired { #[pyo3(item)] partition: String, #[pyo3(item)] graph_walk: String, #[pyo3(item)] node_name: String, @@ -1134,12 +1134,54 @@ pub struct RouteArg { #[pyo3(item)] rids: Vec, #[pyo3(item)] tensors: Vec, #[pyo3(item)] num_tensors: Vec, +} + +/// `RouteInput`. +pub struct RouteArg { + partition: String, + graph_walk: String, + node_name: String, + output_signals: Vec, + rids: Vec, + tensors: Vec, + num_tensors: Vec, /// One output signal whose per-request value is a scalar the worker /// already holds on the host (a sampled token). Its EMIT_TO_CLIENT edge /// then rides inline in one frame per step instead of as a tensor per /// request; `inline_values` are in `rids` order, or empty for none. - #[pyo3(item)] inline_signal: Option, - #[pyo3(item)] inline_values: Vec, + /// Both keys may be absent or None: a caller with no inline signal sends + /// the plain route input. + inline_signal: Option, + inline_values: Vec, +} + +/// An item of a mapping that may be missing or None. +fn optional_item<'py, T: FromPyObject<'py>>( + ob: &Bound<'py, PyAny>, key: &str, +) -> PyResult> { + match ob.get_item(key) { + Ok(v) if v.is_none() => Ok(None), + Ok(v) => v.extract().map(Some), + Err(e) if e.is_instance_of::(ob.py()) => Ok(None), + Err(e) => Err(e), + } +} + +impl<'py> FromPyObject<'py> for RouteArg { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let r: RouteArgRequired = ob.extract()?; + Ok(RouteArg { + partition: r.partition, + graph_walk: r.graph_walk, + node_name: r.node_name, + output_signals: r.output_signals, + rids: r.rids, + tensors: r.tensors, + num_tensors: r.num_tensors, + inline_signal: optional_item(ob, "inline_signal")?, + inline_values: optional_item(ob, "inline_values")?.unwrap_or_default(), + }) + } } /// `RouteOutput`. From 4406b623b0a52ad10ea3836b73ba005a1749f0a5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 04:37:42 -0700 Subject: [PATCH 098/219] test: row store tests skip without the columnar tensor store --- test/modular/test_row_output_store.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/test/modular/test_row_output_store.py b/test/modular/test_row_output_store.py index 78c9d814ae..9029700df6 100644 --- a/test/modular/test_row_output_store.py +++ b/test/modular/test_row_output_store.py @@ -1,5 +1,6 @@ """The row store: a decode step's outputs are row views of one batch clone, stored from one row's description and the stop check's host rows.""" +import pytest import torch from mstar.communication.tensors import SharedMemoryCommunicationManager @@ -19,6 +20,14 @@ def send(self, *args, **kwargs): ) +def _columnar(mgr): + """The row store only exists for the columnar (mstar_rust) tensor store; + without it the manager takes the general path, which is tested elsewhere.""" + if not mgr.tensor_store.has_put_tensor_batch_columns: + pytest.skip("the row store needs the columnar tensor store (mstar_rust)") + return mgr + + def _batch(rids): rows = torch.arange(10, 10 + len(rids), dtype=torch.long).view(len(rids), 1) clone = rows.clone() @@ -44,7 +53,7 @@ def _batch(rids): def test_row_store_matches_the_general_path(tmp_path): rids = [3, 5, 8, 13] outputs, host_rows, cpu_per_rid = _batch(rids) - fast_mgr, slow_mgr = _manager(tmp_path / "a"), _manager(tmp_path / "b") + fast_mgr, slow_mgr = _columnar(_manager(tmp_path / "a")), _manager(tmp_path / "b") fast = fast_mgr.store_row_outputs_batch( rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, node_name="text", graph_walk="decode", host_rows=host_rows, @@ -99,7 +108,7 @@ def test_row_store_keeps_the_general_path_shape_for_missing_signals(tmp_path): rids = [1, 2, 3] outputs, host_rows, cpu_per_rid = _batch(rids) del outputs.per_rid_outputs[2]["text_inputs"] - mgr = _manager(tmp_path) + mgr = _columnar(_manager(tmp_path)) fast = mgr.store_row_outputs_batch( rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, host_rows=host_rows, @@ -113,7 +122,7 @@ def test_row_store_keeps_the_general_path_shape_for_missing_signals(tmp_path): def test_row_store_without_host_rows_stores_no_host_copy(tmp_path): rids = [1, 2] outputs, _, _ = _batch(rids) - mgr = _manager(tmp_path) + mgr = _columnar(_manager(tmp_path)) fast = mgr.store_row_outputs_batch( rids, outputs, SIGNALS, outputs.row_views, outputs.row_request_ids, ) From 4a527fa2adbc242d0e0e6446886cb0603db4a856 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 13:44:06 -0700 Subject: [PATCH 099/219] event wakeup: a public signal for producers --- mstar/communication/event.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/communication/event.py b/mstar/communication/event.py index 352c07954c..82a8b00d51 100644 --- a/mstar/communication/event.py +++ b/mstar/communication/event.py @@ -9,6 +9,10 @@ def __init__(self): def _wake(self, _fut): # runs on whatever thread finished the future os.eventfd_write(self.event, 1) # one syscall, thread-safe, async-signal-safe + def signal(self): + """Wake the poller from any thread (a producer with new work).""" + os.eventfd_write(self.event, 1) + def register_future(self, future: Future): if future.done(): return From 5fea3467cedd4c98912b4e1285fa76bffa4987c7 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 13:44:06 -0700 Subject: [PATCH 100/219] data worker: block on the communicator's poll when idle instead of a 1 ms sleep --- mstar/api_server/data_worker.py | 51 +++++++++++++++++++++++++++++++-- 1 file changed, 49 insertions(+), 2 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 152f2e121e..16768637de 100644 --- a/mstar/api_server/data_worker.py +++ b/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,6 +243,7 @@ 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, @@ -247,6 +262,7 @@ def new_result_token( self.result_tensor_input_queue.put( InlineToken(request_id, value, loop_indices, modality) ) + self._signal() def discard_result_tensors(self, input: ResultTensors): """Ack and drop result tensors for an already-removed request. @@ -255,6 +271,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 +341,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 +403,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 +448,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. @@ -675,10 +701,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( @@ -928,11 +958,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: @@ -1016,7 +1063,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 From 79cc69acb6157af034183a76b58dfddd27dc4765 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 13:44:14 -0700 Subject: [PATCH 101/219] test: the data worker's idle wait and its wake-ups --- test/modular/test_data_worker_wakeup.py | 170 ++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 test/modular/test_data_worker_wakeup.py diff --git a/test/modular/test_data_worker_wakeup.py b/test/modular/test_data_worker_wakeup.py new file mode 100644 index 0000000000..4b4c9e68ba --- /dev/null +++ b/test/modular/test_data_worker_wakeup.py @@ -0,0 +1,170 @@ +"""The data worker thread's idle wait: it blocks on the communicator's poll +and is woken by a producer's signal, a finished tensor read or shutdown, +instead of sleeping 1 ms and sweeping its queues.""" +from __future__ import annotations + +import queue +import select +import threading +import time +from concurrent.futures import Future + +from mstar.api_server.data_worker import InlineToken, PreprocessWorkerThread +from mstar.communication.event import EventWakeup +from mstar.graph.loop_indices import NestedLoopIndices + + +class _Model: + def postprocess(self, output, modality, request_kwargs=None): + return f"<{int(output[0])}>".encode() + + +class _PollingComm: + """A communicator whose idle wait really blocks on the wakeup's fd.""" + + def __init__(self): + self.event = None + self.waits = 0 + + def register_event_for_poll(self, event): + self.event = event + + def wait_for_work(self, timeout_ms): + self.waits += 1 + ready, _, _ = select.select([self.event.fd], [], [], timeout_ms / 1000.0) + if ready: + self.event.drain() + + def get_all_new_messages(self): + return [] + + def send(self, entity, msg): + pass + + +class _BareComm: + def get_all_new_messages(self): + return [] + + def send(self, entity, msg): + pass + + +class _TM: + def __init__(self): + self.futures: list[Future] = [] + + def start_read_tensors(self, request_id, graph_edges, graph_walk=None): + return list(self.futures) + + def get_ready_tensors(self): + return {} + + def has_inflight_reads(self, request_id): + return False + + def cleanup_request(self, request_id): + pass + + def force_cleanup_request(self, request_id): + pass + + def ack_unread_tensors(self, request_id, graph_edges): + pass + + +def _thread(comm, wakeup, idle_wait_ms=2000): + stop = threading.Event() + worker = PreprocessWorkerThread( + in_queue=queue.Queue(), result_tensor_queue=queue.Queue(), out_queue=queue.Queue(), + profile_queue=queue.Queue(), cleanup_request_queue=queue.Queue(), + abort_request_queue=queue.Queue(), reads_done_queue=queue.Queue(), + discard_tensor_queue=queue.Queue(), stop_event=stop, communicator=comm, + tensor_manager=_TM(), model=_Model(), wakeup=wakeup, + ) + worker.idle_wait_ms = idle_wait_ms + t = threading.Thread(target=worker.run, daemon=True) + t.start() + return worker, stop, t + + +def _loop(i): + return NestedLoopIndices( + loop_name_order=["decode_loop"], loop_indices={"decode_loop": i}, wg_fwd_pass_idx=i + 1, + ) + + +def _wait_until(pred, timeout_s): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(0.001) + return pred() + + +def test_a_signalled_producer_ends_the_idle_wait_at_once(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=2000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) # the thread is parked + t0 = time.monotonic() + worker.result_tensor_queue.put(InlineToken("req", 7, _loop(0), "text")) + wakeup.signal() + chunk = worker.out_queue.get(timeout=1.0) + assert time.monotonic() - t0 < 0.5 # far below the 2 s backstop + assert chunk.request_id == "req" and chunk.data == b"<7>" + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + + +def test_shutdown_wakes_a_parked_thread(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=5000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) + t0 = time.monotonic() + stop.set() + wakeup.signal() + t.join(timeout=2.0) + assert not t.is_alive() and time.monotonic() - t0 < 1.0 + + +def test_an_idle_thread_does_not_spin(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=50) + time.sleep(0.3) + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + # ~6 waits of 50 ms in 0.3 s, not ~300 sweeps of 1 ms + assert comm.waits <= 12 + + +def test_a_finished_read_wakes_the_thread(): + comm, wakeup = _PollingComm(), EventWakeup() + worker, stop, t = _thread(comm, wakeup, idle_wait_ms=2000) + assert _wait_until(lambda: comm.waits >= 1, 1.0) + fut: Future = Future() + worker.tensor_manager.futures = [fut] + worker.wakeup.register_futures(worker.tensor_manager.start_read_tensors("r", [])) + waits = comm.waits + time.sleep(0.05) + assert comm.waits == waits # still parked + fut.set_result(None) # the reader thread finishing wakes the poll + assert _wait_until(lambda: comm.waits > waits, 1.0) + stop.set() + wakeup.signal() + t.join(timeout=1.0) + assert not t.is_alive() + + +def test_without_a_wakeup_the_thread_keeps_the_sleep(): + worker, stop, t = _thread(_BareComm(), None) + worker.result_tensor_queue.put(InlineToken("req", 3, _loop(0), "text")) + chunk = worker.out_queue.get(timeout=1.0) + assert chunk.data == b"<3>" + stop.set() + t.join(timeout=1.0) + assert not t.is_alive() From b261d9bf4d1ab1f27404a262e370cb7b944c44d7 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 13:44:15 -0700 Subject: [PATCH 102/219] ci: list the data worker wake-up test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 94984fced4..6c7bb7df87 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -55,3 +55,4 @@ test/modular/test_publish_snapshot_batch.py test/modular/test_sampler_for_cache.py test/modular/test_send_prep_entries.py test/modular/test_kv_prefix_batch.py +test/modular/test_data_worker_wakeup.py From bf6762c5ed0e6624ab5e6355f5242d193e7e12b7 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:07:59 -0700 Subject: [PATCH 103/219] sampler buffers: a slot-addressed last-token master written by the captured sample --- mstar/engine/resources/sampler/utils.py | 59 +++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index c17af7a48d..b1c9d50d55 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1025,6 +1025,10 @@ class CudaGraphableSampler(BaseSampler): seen_tokens_buf: torch.Tensor | None = None # [bs, V] bool # ``None`` for submodules whose ``SamplerSpec`` leaves ``enable_min_p`` off. min_p_buf: torch.Tensor | None = None + # ``[bs]`` per-step row of the slot's ``last_token`` buffer: the captured + # sample writes the (TP-agreed) token here, and ``commit`` scatters it to + # the slot master so the next step can read the token on the device. + last_token_buf: torch.Tensor | None = None tp_group: "CommGroup | None" = None # noqa: F821 # Set during graph capture, and used by the cuda graph runner to determine @@ -1047,6 +1051,8 @@ def sample( ) self.offset_buf += 1 codes = self._broadcast_tokens(codes) + if self.last_token_buf is not None: + self.last_token_buf.copy_(codes) if apply_penalty and self.seen_tokens_buf is not None: self.applied_penalty_in_graph = True # Record the (broadcast, TP-agreed) token in the seen-token buffer so @@ -1266,6 +1272,14 @@ class SamplerBuffers: # ``_scalar_buffers`` (never written from a SamplingConfig) and reset to 0 # on register instead. offset: Buffer + # Per-request last sampled token, by slot: the captured sample writes the + # per-step row (double-buffered, one row per cg slot), ``scatter_last_token`` + # persists it to the master after the replay, and a decode step reads its + # input ids straight off the master (``gather_last_tokens``) instead of a + # per-request tensor routed through the graph. An eager step writes the + # master by request (``write_last_tokens``). Never gathered into the + # per-step row: only the sample writes it. + last_token: Buffer # TP communicator for the submodule that owns these buffers. Passed # through ``slice_for_bs`` into every per-step ``CudaGraphableSampler`` @@ -1380,6 +1394,7 @@ def mk(dtype: torch.dtype, default: float) -> HostBuffer: seed=mk(torch.long, 0), rep_penalty=mk(torch.float32, 1.0), offset=Buffer.allocate(max_batch_size, cap, device, torch.long, 0, 1), + last_token=Buffer.allocate(max_batch_size, cap, device, torch.long, 0, cg_slots), tp_group=tp_group, seen_tokens=seen_tokens, min_p=mk(torch.float32, 0.0) if enable_min_p else None, @@ -1405,6 +1420,7 @@ def slice_for_bs(self, bs: int, cg_slot: int = 0) -> dict[str, Any]: "rep_penalty_buf": self.rep_penalty.slot_view(cg_slot, bs), "seen_tokens_buf": self.seen_tokens.slot_view(cg_slot, bs) if self.seen_tokens is not None else None, "min_p_buf": self.min_p.slot_view(cg_slot, bs) if self.min_p is not None else None, + "last_token_buf": self.last_token.slot_view(cg_slot, bs), "tp_group": self.tp_group, } @@ -1448,6 +1464,7 @@ def _grow_master(self, new_capacity: int) -> None: for buf in self._scalar_buffers(): buf.grow_master(new_capacity) self.offset.grow_master(new_capacity) + self.last_token.grow_master(new_capacity) if self.seen_tokens is not None: self.seen_tokens.grow_master(new_capacity) self._free_slots.extend(range(self._master_capacity, new_capacity)) @@ -1669,3 +1686,45 @@ def scatter_offset(self, cg_slot: int = 0) -> None: self.offset.scatter( self._slot_idx_gpu[cg_slot], self._last_real_bs[cg_slot], cg_slot ) + + # ------------------------------------------------------------------ + # Last sampled token: written by the sample, read by the next step + # ------------------------------------------------------------------ + def scatter_last_token(self, cg_slot: int = 0) -> None: + """Persist the tokens the captured sample wrote into ``cg_slot``'s + per-step row to their slot masters. Same contract as ``scatter_offset``: + after the replay, real rows only.""" + self.last_token.scatter( + self._slot_idx_gpu[cg_slot], self._last_real_bs[cg_slot], cg_slot + ) + + def gather_last_tokens(self, cg_slot: int, padded_bs: int) -> torch.Tensor: + """``[padded_bs]`` last tokens of the batch staged on ``cg_slot``, in + its row order (padding rows read slot 0, like every other gather). + A fresh tensor on the current stream; the runner stages it into the + slot's static ``input_ids``. Call after ``gather_dynamic`` uploaded the + slot's index row for this step.""" + return self.last_token.master.index_select( + 0, self._slot_idx_gpu[cg_slot, :padded_bs] + ) + + def _slots_of(self, request_ids: list[str]) -> torch.Tensor: + slots = [self._rid_to_slot[rid] for rid in request_ids] + return torch.tensor(slots, dtype=torch.long, device=self.last_token.master.device) + + def write_last_tokens(self, request_ids: list[str], tokens: torch.Tensor) -> None: + """Eager-step write of ``tokens`` (``[len(request_ids)]``) to the + requests' master rows. Every request must be registered.""" + if not request_ids: + return + self.last_token.master.index_copy_( + 0, self._slots_of(request_ids), tokens.to(self.last_token.master.dtype) + ) + + def last_tokens_for(self, request_ids: list[str]) -> torch.Tensor: + """``[len(request_ids)]`` last tokens of registered requests, for an + eager step that has no staged slot row.""" + return self.last_token.master.index_select(0, self._slots_of(request_ids)) + + def has_slot(self, rid: str) -> bool: + return rid in self._rid_to_slot From 9e63ef05d6a14505a05395469973c950edf448bf Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:07:59 -0700 Subject: [PATCH 104/219] sampler resource: scatter the step's tokens at commit, read them back by slot or by request --- mstar/engine/resources/sampler/resource.py | 35 +++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index d072523d87..fbed9c20ea 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -70,6 +70,9 @@ def __init__( self._preplan_cg_sampler: CudaGraphableSampler | None = None self._preplan_key = None self._preplanned = False + # (cg slot, padded bs) of the step planned last, for the last-token + # read in ``loopback_tokens``; None on an eager step + self._step_slot: tuple[int, int] | None = None # rid -> the prompt tokens a cache hit kept out of this step's inputs self._cached_prefix: dict[str, torch.Tensor] = {} @@ -259,12 +262,14 @@ def plan(self, step: SamplerStep, ctx: StepContext): self._cg_sampler = self._preplan_cg_sampler self._preplan_cg_sampler = None self._preplanned = False + self._step_slot = (ctx.slot_lease.slot, ctx.slot_lease.bucket.bs) return # invalidated on the inline path here (not in commit, which now runs # before output collection); a preplan must leave the in-flight one be if not ctx.is_preplan: self._cg_sampler = None + self._step_slot = None lease = ctx.slot_lease if self._cg_buffers is None or lease is None: return @@ -286,6 +291,7 @@ def plan(self, step: SamplerStep, ctx: StepContext): # fresh inline (capture / no preplan): gather the per-step state too self._gather_dynamic(ctx, lease) self._cg_sampler = sampler + self._step_slot = (cg_slot, padded_bs) def _gather_dynamic(self, ctx: StepContext, lease: SlotLease): """Gather the per-step RNG offset + seen-token mask into this step's @@ -311,6 +317,7 @@ def commit(self, step: SamplerStep, ctx: StepContext): if self._cg_buffers is None or self._cg_sampler is None: return self._cg_buffers.scatter_offset(ctx.slot_lease.slot) + self._cg_buffers.scatter_last_token(ctx.slot_lease.slot) # Skipped when nothing read the mask this step: there is then nothing to # copy back, and the rows go stale only for requests at penalty 1.0, # which never read them. See `_penalty_live` for why that stays sound. @@ -341,4 +348,30 @@ def sample( request_ids, logits, apply_penalty=self._apply_penalty_this_step ) - return self._sampler.sample(request_ids, logits) + tokens = self._sampler.sample(request_ids, logits) + # An eager step has no per-step row to scatter from: write the masters + # by request, so the next step can still read its token on the device. + bufs = self._cg_buffers + if bufs is not None and all(bufs.has_slot(rid) for rid in request_ids): + bufs.write_last_tokens(request_ids, tokens) + return tokens + + @property + def has_slot_masters(self) -> bool: + """Whether the last sampled token of every request is kept on the + device by slot (the graph-buffer path): what the device-side + loop-back needs. False on an eager-only node.""" + return self._cg_buffers is not None + + def loopback_tokens(self, request_ids: list[str]) -> torch.Tensor: + """The last sampled token of each row of this step, as its input ids, + read on the device. Under a slot lease the planned slot's index row + orders them (padding rows read slot 0); an eager step looks its + requests up by slot. Call after ``plan``.""" + bufs = self._cg_buffers + if bufs is None: + raise RuntimeError("device loop-back needs the sampler's slot masters") + if self._cg_sampler is not None and self._step_slot is not None: + cg_slot, padded_bs = self._step_slot + return bufs.gather_last_tokens(cg_slot, padded_bs) + return bufs.last_tokens_for(request_ids) From bb28b086f21f628fc06744eacff83ef902705595 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:07:59 -0700 Subject: [PATCH 105/219] submodule base: the device loop-back hook and its knob --- mstar/model/submodule_base.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index d2d22d668d..8fd13212b1 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Mapping, Sequence @@ -258,6 +259,12 @@ def _split_pos_ids( return pos_ids[..., start:end] +def device_loopback_enabled() -> bool: + """``MSTAR_DEVICE_LOOPBACK`` (default 1): whether a node may keep its + loop-back token on the device instead of routing a tensor per request.""" + return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" + + @dataclass class ARNodeInputs(NodeInputs): """ @@ -692,6 +699,17 @@ def step_is_input_free(self, graph_walk: str) -> bool: del graph_walk return False + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: + """Loop-back input signals of this walk whose per-row value the node + keeps on the device itself (a sampler's slot master), so the worker + routes them with no tensor: the next step reads the value in + ``preprocess`` and every fallback batch arrives with the signal + empty. Only sound when ``prepare_inputs`` accepts the empty signal on + this walk. Empty by default; ``MSTAR_DEVICE_LOOPBACK=0`` keeps every + node on the tensor path.""" + del graph_walk + return frozenset() + def inline_client_signals(self, graph_walk: str) -> dict[str, str]: """Output signals whose client-facing value is one scalar per row that the stop check already copies to the host, as signal name -> the name From 970c8cec1da2271c18cd9f00fd960add65f79089 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 106/219] qwen3.5: decode input ids off the sampler's slot master --- mstar/model/qwen3_5/submodules.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 7b2dbaf006..581eb4f4a2 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -69,6 +69,7 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + device_loopback_enabled, ) logger = logging.getLogger(__name__) @@ -276,9 +277,24 @@ def prepare_inputs( if graph_walk == "prefill_vision": return self._vision_inputs(fwd_info, inputs) + if graph_walk == "decode" and self._device_loopback(): + # one token per row, read off the sampler's slot master in + # `preprocess`; the routed signal is empty (or ignored) here + return ARNodeInputs(input_seq_len=1) input_ids = inputs["text_inputs"][0] return ARNodeInputs(input_seq_len=input_ids.shape[0], input_ids=input_ids) + def _device_loopback(self) -> bool: + return ( + device_loopback_enabled() + and self.node_resources[SAMPLER].has_slot_masters + ) + + def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: + if graph_walk == "decode" and self._device_loopback(): + return frozenset({"text_inputs"}) + return frozenset() + def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: """``[3, total_tokens]`` for the step, in packed request order. @@ -317,7 +333,12 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: out: dict[str, torch.Tensor | Any] = {} - if inputs[0].input_ids is not None: + if inputs[0].input_ids is None and inputs[0].input_embeds is None: + # device loop-back: every row of a decode step is its request's + # last sampled token, read by slot (padding rows read slot 0) + sampler: SamplerResource = engine_inputs.resources[SAMPLER] + out["input_ids"] = sampler.loopback_tokens(engine_inputs.request_ids) + elif inputs[0].input_ids is not None: out["input_ids"] = torch.cat([inp.input_ids for inp in inputs], dim=0) else: out["input_embeds"] = torch.cat( From 954bace71cd910786b38368734ee9887afb04bf0 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 107/219] engine: expose the node's device loop-back signals --- mstar/engine/engine.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 4e54522116..1dc1d31c95 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -796,6 +796,16 @@ def inline_client_signals( return {} return mgmt.submodule.inline_client_signals(graph_walk) + def device_loopback_signals( + self, node_name: str, graph_walk: str, + ) -> frozenset[str]: + """The node's loop-back signals it keeps on the device this walk, so + the worker routes them without a tensor (see the submodule hook).""" + mgmt = self._submodules.get(node_name) + if mgmt is None: + return frozenset() + return mgmt.submodule.device_loopback_signals(graph_walk) + def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], host_rows: HostRows | None = None, From cf9b4f8b2c5cfb572e54d9a13e458f243b0367c9 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 108/219] worker: no store, hold or tensors for a step kept on the device --- mstar/worker/worker.py | 39 ++++++++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index ce8154cae2..65085892d8 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -22,6 +22,7 @@ from mstar.communication.tensors import ( LocalTransferEngine, NameToTensorList, + StoredOutputs, create_tensor_communication_manager, ) from mstar.conductor.request_info import CurrentForwardPassInfo @@ -298,6 +299,8 @@ def __init__( self.worker_id = worker_id self.device = device self.enable_nvtx = nvtx_enabled(enable_nvtx, device) + # (node, walk) -> the loop-back signals the node keeps on the device + self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -2974,8 +2977,22 @@ def _pp_stage(name: str) -> None: # reuses; safe because the sends below are their last reader. # Row views of one batch clone (a decode step) are described from one # row; anything else goes through the general per-tensor path. + # The client copy first: whether it rides inline decides whether the + # tensors are needed at all. + inline_signal, inline_values = self._inline_emit_values( + engine, batch_N, host_rows, rids, signals, + ) stored = None - if outputs.row_views and outputs.row_request_ids is not None: + if inline_signal is not None and self._all_on_device( + engine, batch_N.node_name, batch_N.graph_walk, signals, + ): + # Device loop-back: the node keeps every routed signal's value on + # the device (the sampler's slot master) and the client copy rides + # inline, so nothing is stored, held or freed. The route still + # runs per request for the loop counters and the fallbacks; it + # carries zero tensors per (request, signal). + stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) + elif outputs.row_views and outputs.row_request_ids is not None: stored = self.tensor_manager.store_row_outputs_batch( rids, outputs, signals, outputs.row_views, outputs.row_request_ids, node_name=batch_N.node_name, @@ -2999,7 +3016,8 @@ def _pp_stage(name: str) -> None: # one per tensor -- with a Rust bookkeeper each is a boundary crossing, # and a 128-request batch has hundreds of them. The count is uniform, # so it crosses as a scalar rather than a list built per batch. - self.tensor_manager.increment_ref_batch_uniform(flat_uuids, 1) + if flat_uuids: + self.tensor_manager.increment_ref_batch_uniform(flat_uuids, 1) if self._phase_period: self._phase_record( "worker.postprocess.store_tensors", @@ -3008,9 +3026,6 @@ def _pp_stage(name: str) -> None: # The graph runtime's own share of postprocess: the routing call. _t_route = _time.perf_counter() if self._phase_period else 0.0 - inline_signal, inline_values = self._inline_emit_values( - engine, batch_N, host_rows, rids, signals, - ) route_output = self._graph_runtime.complete_and_route_batch( RouteInput( partition=batch_N.partition, @@ -3196,6 +3211,20 @@ def _pp_stage(name: str) -> None: if self.enable_nvtx: range_pop(synchronize=False) + def _all_on_device(self, engine, node_name: str, graph_walk: str, signals) -> bool: + """Whether every output signal of this step is one the node keeps on + the device (cached per node and walk: it depends on the node's + resources, not on the step).""" + if not signals: + return False + key = (node_name, graph_walk) + on_device = self._device_loopback.get(key) + if on_device is None: + on_device = self._device_loopback[key] = frozenset( + engine.device_loopback_signals(node_name, graph_walk) + ) + return all(s in on_device for s in signals) + def _inline_emit_values( self, engine, batch_N: PendingBatch, host_rows: "HostRows | None", rids: list[int], signals, From 05ebf8e7fcaa5d5d61c83d0da4064b415814ae94 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 109/219] docs: MSTAR_DEVICE_LOOPBACK --- docs/environment_variables.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index e3a821e4fc..49a680ee0a 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -332,6 +332,15 @@ Worker scheduling milliseconds lets the arrivals of that window prefill together, which is worth decode throughput at high concurrency and costs up to that much time to first token. + * - ``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_RUST_SEND_TIMING`` - unset - ``1`` makes the Rust runtime print, every 500 calls, how long the From 5ea711fe2402a6c1d77c61ac4ae9332547e217a0 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 110/219] test: the sampler's last-token master --- test/modular/test_sampler_last_token.py | 61 +++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 test/modular/test_sampler_last_token.py diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py new file mode 100644 index 0000000000..d2cb346744 --- /dev/null +++ b/test/modular/test_sampler_last_token.py @@ -0,0 +1,61 @@ +"""The sampler's last-token master: the device-side loop-back's store. The +captured sample writes a per-step row, commit scatters it to the slot master, +the next step gathers its input ids by slot; an eager step writes and reads +the master by request.""" +import torch + +from mstar.engine.resources.sampler.utils import SamplerBuffers + + +def _buffers(max_bs=4, cg_slots=2): + return SamplerBuffers.allocate( + max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, + ) + + +def test_eager_write_and_read_go_by_request(): + bufs = _buffers() + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b"], torch.tensor([11, 12])) + assert bufs.last_tokens_for(["b", "a"]).tolist() == [12, 11] + bufs.write_last_tokens(["c"], torch.tensor([13], dtype=torch.int32)) + assert bufs.last_tokens_for(["c", "a"]).tolist() == [13, 11] + assert bufs.has_slot("a") and not bufs.has_slot("zz") + + +def test_the_captured_row_scatters_to_the_slot_master_and_gathers_back(): + bufs = _buffers(max_bs=4, cg_slots=2) + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b", "c"], torch.tensor([1, 2, 3])) + # a step of [b, a] padded to 4 on cg slot 1: static then dynamic gather + bufs.gather_static(["b", "a"], 4, 1) + bufs.gather_dynamic(["b", "a"], 4, 1) + # what the captured sample does: write the step's tokens into the slot row + bufs.slice_for_bs(4, 1)["last_token_buf"].copy_(torch.tensor([22, 21, 99, 99])) + bufs.scatter_last_token(1) + # real rows reached their masters, padding rows did not touch slot 0's + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [21, 22, 3] + # the next step's input ids, in the slot row's order; padding rows read + # slot 0 (whatever request holds it, or the default when it is free) + slot0_tok = bufs.last_token.master[0].item() + assert bufs.gather_last_tokens(1, 4).tolist() == [22, 21, slot0_tok, slot0_tok] + # the other cg slot is untouched by this step + bufs.gather_static(["c"], 4, 0) + bufs.gather_dynamic(["c"], 4, 0) + assert bufs.gather_last_tokens(0, 4).tolist()[0] == 3 + + +def test_growth_and_slot_reuse_keep_the_tokens(): + bufs = _buffers(max_bs=2, cg_slots=1) + bufs.register_request("a") + bufs.register_request("b") + bufs.write_last_tokens(["a", "b"], torch.tensor([5, 6])) + bufs.register_request("c") # beyond the capacity of 2: the masters grow + bufs.write_last_tokens(["c"], torch.tensor([7])) + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [5, 6, 7] + bufs.unregister_request("a") + bufs.register_request("d") # reuses a's slot; its first step writes it + bufs.write_last_tokens(["d"], torch.tensor([8])) + assert bufs.last_tokens_for(["d", "b", "c"]).tolist() == [8, 6, 7] From 21f920a390bb6ec7102ee81b3bc0a2749ef7412e Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 111/219] test: qwen3.5 decode inputs from the sampler master --- test/modular/test_qwen3_5_device_loopback.py | 80 ++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 test/modular/test_qwen3_5_device_loopback.py diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py new file mode 100644 index 0000000000..dedab1d856 --- /dev/null +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -0,0 +1,80 @@ +"""Qwen3.5 decode with the device-side loop-back: the step's input ids come +from the sampler's slot master, the routed signal carries no tensor, and the +knob or a sampler without masters keeps the tensor path.""" +from types import SimpleNamespace + +import pytest +import torch + +from mstar.model.qwen3_5.config import SAMPLER +from mstar.model.qwen3_5.submodules import LLMSubmodule + + +class _Sampler: + def __init__(self, masters=True, tokens=None): + self.has_slot_masters = masters + self.tokens = tokens + self.asked = None + + def loopback_tokens(self, request_ids): + self.asked = list(request_ids) + return self.tokens + + +def _sub(sampler): + sub = LLMSubmodule.__new__(LLMSubmodule) + sub.node_resources = {SAMPLER: sampler} + return sub + + +def _prepare(sub, walk, inputs): + return LLMSubmodule.prepare_inputs(sub, graph_walk=walk, fwd_info=None, inputs=inputs) + + +def test_decode_takes_no_input_ids_when_the_sampler_keeps_them(): + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {}) + assert out.input_ids is None and out.input_embeds is None and out.input_seq_len == 1 + # a routed tensor is ignored on this path: the master is the source + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}) + assert out.input_ids is None + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset({"text_inputs"}) + assert LLMSubmodule.device_loopback_signals(sub, "prefill_text") == frozenset() + + +def test_without_masters_or_with_the_knob_off_the_tensor_path_stays(monkeypatch): + sub = _sub(_Sampler(masters=False)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}) + assert out.input_ids.tolist() == [7] and out.input_seq_len == 1 + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset() + with pytest.raises(KeyError): + _prepare(sub, "decode", {}) # the tensor path needs the signal + monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([9])]}) + assert out.input_ids.tolist() == [9] + assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset() + + +def test_preprocess_reads_the_step_tokens_off_the_sampler(): + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 0, 0])) + sub = _sub(sampler) + rows = [_prepare(sub, "decode", {}) for _ in range(4)] + engine_inputs = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2", "pad1", "pad2"], + ) + seen = {} + + def _position_ids_3d(inputs): + seen["n"] = len(inputs) + return torch.zeros(3, len(inputs)) + + sub.model = SimpleNamespace(model=SimpleNamespace( + build_cos_sin=lambda pos, dtype: (pos, pos), + embed_tokens=SimpleNamespace(weight=torch.zeros(1)), + )) + sub._position_ids_3d = _position_ids_3d + out = LLMSubmodule.preprocess(sub, "decode", engine_inputs, rows) + assert out["input_ids"].tolist() == [3, 4, 0, 0] + assert sampler.asked == ["r1", "r2", "pad1", "pad2"] + assert "input_embeds" not in out and seen["n"] == 4 From 4f0fadb6244d071f6c607c63c2bdf17423293511 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 112/219] test: the worker's device loop-back decision --- test/modular/test_worker_device_loopback.py | 44 +++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 test/modular/test_worker_device_loopback.py diff --git a/test/modular/test_worker_device_loopback.py b/test/modular/test_worker_device_loopback.py new file mode 100644 index 0000000000..6165b3c966 --- /dev/null +++ b/test/modular/test_worker_device_loopback.py @@ -0,0 +1,44 @@ +"""The worker's side of the device loop-back: a step whose every routed signal +the node keeps on the device, with the client copy inline, stores and holds +nothing and routes zero tensors per (request, signal).""" +from types import SimpleNamespace + +from mstar.communication.tensors import StoredOutputs +from mstar.worker.worker import Worker + + +class _Engine: + def __init__(self, on_device): + self.on_device = on_device + self.calls = 0 + + def device_loopback_signals(self, node, walk): + self.calls += 1 + return self.on_device.get((node, walk), frozenset()) + + +def _worker(): + w = Worker.__new__(Worker) + w._device_loopback = {} + return w + + +def test_all_on_device_is_per_step_signals_and_cached_per_node_and_walk(): + engine = _Engine({("LLM", "decode"): frozenset({"text_inputs"})}) + w = _worker() + assert w._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + assert w._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + assert engine.calls == 1, "the node's answer is cached per (node, walk)" + # a step that also routes a signal the node does not keep stays on tensors + assert not w._all_on_device(engine, "LLM", "decode", ["text_inputs", "kv"]) + assert not w._all_on_device(engine, "LLM", "prefill_text", ["new_token"]) + assert not w._all_on_device(engine, "LLM", "decode", []) + assert engine.calls == 2 + + +def test_the_zero_tensor_store_shape_matches_the_route_layout(): + rids, signals = [3, 5, 8], ["text_inputs"] + stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) + assert stored.flat_uuids == [] and stored.flat_rids == [] and stored.signal_idxs == [] + assert stored.num_tensors == [0, 0, 0] + assert SimpleNamespace(**stored._asdict()).num_tensors == [0, 0, 0] if hasattr(stored, "_asdict") else True From d57c15d39051d832b3ccb09cbb15f98413847737 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 113/219] test: a loop-back routed with no tensor keeps the loop going --- test/rust/test_rust_graph_runtime.py | 35 ++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/test/rust/test_rust_graph_runtime.py b/test/rust/test_rust_graph_runtime.py index 1faec35751..569b65f59d 100644 --- a/test/rust/test_rust_graph_runtime.py +++ b/test/rust/test_rust_graph_runtime.py @@ -1025,3 +1025,38 @@ def test_inline_values_must_cover_every_rid_or_are_ignored(): )) assert 201 in out.register_uuids, "no values, so the tensor path" rt.send_outputs(_send_input(rid, out.completion_id)) + + +def test_a_loop_back_routed_with_no_tensor_keeps_the_loop_going(): + """The device-side loop-back: the worker routes the decode signal with + zero tensors per request (the token stays on the device) and the client + copy rides inline. The node stays ready, the loop counter advances, the + next pop reports the signal with no uuids, and cleanup frees nothing.""" + rt = _client_loop_runtime() + rid = _admit(rt) + rt._communicator = _Recorder() + _reach_dec(rt, rid, 300) + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[], num_tensors=[0], inline_signal="token", inline_values=[7], + )) + assert out.register_uuids == [], "nothing to stage for the client" + rt.send_outputs(_send_input(rid, out.completion_id)) + popped = rt.pop_rids("dec", WALK, [rid]) + assert popped is not None, "the loop-back with no tensor still fed the node" + inputs = popped.input_edges.to_input_tensors(lambda u: (_ for _ in ()).throw(KeyError(u))) + assert inputs.by_rid[rid] == {"token": []} + before = dict(rt.get_dynamic_loop_iters([rid], "default").values[0]) + # another zero-tensor step: still ready, nothing freed by the cleanup, the + # loop counter moved + out = rt.complete_and_route_batch(RouteInput( + partition="default", graph_walk=WALK, node_name="dec", + output_signals=["token"], wg_ids=ParallelList([rid], [WG_ID]), + tensors=[], num_tensors=[0], inline_signal="token", inline_values=[8], + )) + assert out.freed_inputs.uuids == [] + rt.send_outputs(_send_input(rid, out.completion_id)) + assert rt.pop_rids("dec", WALK, [rid]) is not None + after = dict(rt.get_dynamic_loop_iters([rid], "default").values[0]) + assert after["loop"] == before["loop"] + 1 From 1a15e53afdbb33850481e31cf41255c8d4b8e473 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:08:00 -0700 Subject: [PATCH 114/219] ci: list the device loop-back tests --- test/cpu-core.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 6c7bb7df87..766759e150 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -56,3 +56,6 @@ test/modular/test_sampler_for_cache.py test/modular/test_send_prep_entries.py test/modular/test_kv_prefix_batch.py test/modular/test_data_worker_wakeup.py +test/modular/test_sampler_last_token.py +test/modular/test_qwen3_5_device_loopback.py +test/modular/test_worker_device_loopback.py From 874940e03bafdce42eff2e1e86769dd3a1ff23f5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:42:54 -0700 Subject: [PATCH 115/219] worker: early speculation on TP leaders, committing only with a head in hand --- mstar/worker/worker.py | 240 +++++++++++++++++++++++------------------ 1 file changed, 135 insertions(+), 105 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 65085892d8..c5ca557ae0 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -149,6 +149,18 @@ def _parse_tp_async_sched(raw: str) -> tuple[bool, frozenset[str] | None]: return True, frozenset(n.strip() for n in raw.split(",") if n.strip()) +@dataclass +class _SpecFairness: + """The decode loop's fairness settings and clock for ``_build_speculation``: + whether to peek for other ready work, how long to hold the yield once it + is seen, the consecutive-step cap, and when other work first became + ready (None while nothing else waits).""" + peek: bool + hold_s: float + max_consecutive: int + ready_since: float | None = None + + @dataclass class PendingBatch: batch: ScheduledBatch @@ -3599,6 +3611,112 @@ def _fail_requests(self, errors: dict[int, str]) -> None: # Note: we do not cleanup the request right now; we wait for the conductor # to officially send a removal message + def _build_speculation( + self, pending: PendingBatch | None, consecutive: int, + fair: "_SpecFairness", decide: bool = True, + ) -> tuple[Speculation | None, tuple[str, str] | None]: + """Speculate the step after ``pending`` (N+1 from N), or a + yield-away batch when fairness says so. Returns the speculation + and the (node, walk) to exclude on the non-speculative path. + + ``decide=False`` (the early build on a parallel node) makes a + build that finds no head non-committal: no marker to the + followers, no yield-away batch; the caller's next full call + takes the decision for this step instead.""" + speculation = None + yield_away_from_target = None + batch = None + admit_refused = ( + pending is not None + and pending.node_batch.admit_error is not None + ) + + if pending is not None and self._can_speculate(pending.batch): + # Fairness check (peek-based, replaces the old iter- + # counter cap): only break the spec chain when there's + # another (node, walk) actually ready to schedule on + # this worker. On single-walk workers (Orpheus LLM, + # Orpheus SNAC) this returns False and we always speculate. + other_ready = ( + fair.peek + and consecutive >= 1 + and self.scheduler.has_ready_excluding( + self.request_state, + (pending.node_name, pending.graph_walk), + ) + ) + if other_ready: + now = _time.perf_counter() + if fair.ready_since is None: + fair.ready_since = now + must_yield_for_fairness = ( + now - fair.ready_since >= fair.hold_s + ) + else: + fair.ready_since = None + must_yield_for_fairness = False + must_yield_away = ( + consecutive >= fair.max_consecutive + or must_yield_for_fairness + ) + if must_yield_away: + fair.ready_since = None + if not must_yield_away and not admit_refused: + if self.enable_nvtx: + range_push("worker.speculate", synchronize=False) + _t0 = _time.perf_counter() if self._phase_period else 0.0 + speculation = self._try_speculate_next(pending) + if self._phase_period: + self._phase_record("speculate", _time.perf_counter() - _t0) + if self.enable_nvtx: + range_pop(synchronize=False) + if speculation is not None: + # Broadcast the head now, during forward N, so + # followers build it during theirs (-1: not parallel). + speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( + speculation.node_batch, + speculative=True, spec_from_seq=pending.tp_seq, + ) + if speculation is None and not decide: + return None, None + if self._tp_lead_needs_marker(pending, speculation): + self._broadcast_tp_nospec(pending) + # ``yield_away_from_target`` stays None when the admit + # refused, so the non-speculative path below schedules + # without the fairness exclusion — it should be free to pick + # up the batch ``_handle_admit_failure`` just pushed back. + if speculation is None and not admit_refused: + yield_away_from_target = ( + pending.node_name, + pending.graph_walk, + ) if must_yield_away else None + with self._span("worker.schedule_yield_away"): + batch = self.scheduler.get_next_batch( + self.request_state, + exclude_target=yield_away_from_target, + ) + if batch is not None: + node_batch = self._build_executing_batch(batch) + batch_partition = self.request_state.get_partition_for_node(batch.node_name) + logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") + speculation = Speculation( + scheduled_batch=batch, + node_batch=node_batch, + consumed_edges=set(), + continuing_rids=set(), # n/a + partition=batch_partition, + is_new_iter=False, + is_same_node=False, + is_yield_away=True + ) + + # A leader stamps the seq it sends; a follower's + # batch keeps the one it came off the FIFO with. + ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) + speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq + + return speculation, yield_away_from_target + def run(self) -> None: switch_interval = os.environ.get("MSTAR_PY_SWITCH_INTERVAL_SEC", "") if switch_interval: @@ -3725,115 +3843,24 @@ def run(self) -> None: # that at high concurrency the requests arriving over that window # prefill in one step rather than one step each. 0 yields at once. spec_yield_hold_s = float(os.environ.get("MSTAR_SPEC_YIELD_HOLD_MS", "0")) / 1000.0 - fairness_ready_since: list[float | None] = [None] + fair = _SpecFairness( + peek=spec_peek_for_fairness, hold_s=spec_yield_hold_s, + max_consecutive=max_consecutive_spec, + ) consecutive_spec_steps = 0 yield_away_from_target: tuple[str, str] | None = None # Build N+2 while N+1 runs, between the two halves of N's # post-processing: its pre-plan then overlaps the routing half rather - # than the next launch. TP1 nodes only; a parallel node keeps the head - # broadcast in the order the followers settle it. + # than the next launch. On a parallel (TP) node the early build only + # commits when it finds a head: the followers settle each step on + # exactly one decision, so an early build that finds nothing sends no + # marker and schedules no yield-away, and the top of the next + # iteration decides as it always did. early_spec = os.environ.get("MSTAR_EARLY_SPEC", "1") == "1" next_speculation: Speculation | None = None next_yield_away: tuple[str, str] | None = None from mstar.utils.profiler import range_pop, range_push - def _build_speculation( - pending: PendingBatch | None, consecutive: int, - ) -> tuple[Speculation | None, tuple[str, str] | None]: - """Speculate the step after ``pending`` (N+1 from N), or a - yield-away batch when fairness says so. Returns the speculation - and the (node, walk) to exclude on the non-speculative path.""" - speculation = None - yield_away_from_target = None - batch = None - admit_refused = ( - pending is not None - and pending.node_batch.admit_error is not None - ) - - if pending is not None and self._can_speculate(pending.batch): - # Fairness check (peek-based, replaces the old iter- - # counter cap): only break the spec chain when there's - # another (node, walk) actually ready to schedule on - # this worker. On single-walk workers (Orpheus LLM, - # Orpheus SNAC) this returns False and we always speculate. - other_ready = ( - spec_peek_for_fairness - and consecutive >= 1 - and self.scheduler.has_ready_excluding( - self.request_state, - (pending.node_name, pending.graph_walk), - ) - ) - if other_ready: - now = _time.perf_counter() - if fairness_ready_since[0] is None: - fairness_ready_since[0] = now - must_yield_for_fairness = ( - now - fairness_ready_since[0] >= spec_yield_hold_s - ) - else: - fairness_ready_since[0] = None - must_yield_for_fairness = False - must_yield_away = ( - consecutive >= max_consecutive_spec - or must_yield_for_fairness - ) - if must_yield_away: - fairness_ready_since[0] = None - if not must_yield_away and not admit_refused: - if self.enable_nvtx: - range_push("worker.speculate", synchronize=False) - _t0 = _time.perf_counter() if phase_period else 0.0 - speculation = self._try_speculate_next(pending) - if phase_period: - _phase_record("speculate", _time.perf_counter() - _t0) - if self.enable_nvtx: - range_pop(synchronize=False) - if speculation is not None: - # Broadcast the head now, during forward N, so - # followers build it during theirs (-1: not parallel). - speculation.tp_seq = self.maybe_send_zmq_to_tp_followers( - speculation.node_batch, - speculative=True, spec_from_seq=pending.tp_seq, - ) - if self._tp_lead_needs_marker(pending, speculation): - self._broadcast_tp_nospec(pending) - # ``yield_away_from_target`` stays None when the admit - # refused, so the non-speculative path below schedules - # without the fairness exclusion — it should be free to pick - # up the batch ``_handle_admit_failure`` just pushed back. - if speculation is None and not admit_refused: - yield_away_from_target = ( - pending.node_name, - pending.graph_walk, - ) if must_yield_away else None - with self._span("worker.schedule_yield_away"): - batch = self.scheduler.get_next_batch( - self.request_state, - exclude_target=yield_away_from_target, - ) - if batch is not None: - node_batch = self._build_executing_batch(batch) - batch_partition = self.request_state.get_partition_for_node(batch.node_name) - logger.debug(f"Yield away: {batch.node_name} {node_batch.request_ids}") - speculation = Speculation( - scheduled_batch=batch, - node_batch=node_batch, - consumed_edges=set(), - continuing_rids=set(), # n/a - partition=batch_partition, - is_new_iter=False, - is_same_node=False, - is_yield_away=True - ) - - # A leader stamps the seq it sends; a follower's - # batch keeps the one it came off the FIFO with. - ya_seq = self.maybe_send_zmq_to_tp_followers(node_batch) - speculation.tp_seq = ya_seq if ya_seq >= 0 else batch.tp_seq - - return speculation, yield_away_from_target def _arm_speculation( spec: Speculation, pending_for: PendingBatch | None = None, @@ -3961,8 +3988,8 @@ def _phase_flush() -> None: next_speculation = next_yield_away = None self._armed_spec_rids = set() else: - speculation, yield_away_from_target = _build_speculation( - pending, consecutive_spec_steps, + speculation, yield_away_from_target = self._build_speculation( + pending, consecutive_spec_steps, fair, ) if speculation is not None: _arm_speculation(speculation) @@ -4148,7 +4175,6 @@ def _maybe_clear_spec(): if ( early_spec and pp_state is not None and spec_pending is not None - and spec_pending.node_name not in self.parallel_nodes ): # N's stops are decided and N+1 is launched: # build N+2 now and hand its pre-plan to the @@ -4160,9 +4186,13 @@ def _maybe_clear_spec(): self._await_admit_settled(spec_pending) if spec_pending.node_batch.admit_error is None: next_speculation, next_yield_away = ( - _build_speculation( + self._build_speculation( spec_pending, - consecutive_spec_steps + 1, + consecutive_spec_steps + 1, fair, + # a TP leader commits early + # only with a head in hand + decide=spec_pending.node_name + not in self.parallel_nodes, ) ) if next_speculation is not None: From 0abe28ba96695c944b898cf34b59f332d6d9d1f9 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:42:54 -0700 Subject: [PATCH 116/219] docs: MSTAR_EARLY_SPEC on tensor-parallel nodes --- docs/environment_variables.rst | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 49a680ee0a..d88aa6760f 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -309,8 +309,11 @@ Worker scheduling work the main thread does anyway instead of the next launch. ``0`` builds each speculation at the top of the next iteration, after the whole post-processing, which left the plan thread idle during it. - Non-parallel (TP1) nodes only; a tensor-parallel node keeps the old - order so the followers see heads in the order they settle them. + On a tensor-parallel node the leader's early build commits only when + it finds a head (it broadcasts it at once); when it finds nothing it + sends no marker and schedules no yield-away, and the top of the next + iteration decides, so the followers still settle every step on + exactly one decision. * - ``MSTAR_LAZY_PUBLISH`` - ``1`` - Publish resource state (KV pages, positions) only for the requests From b794f72c11df503b69e1e7e5d8bcd37301978dbc Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:42:54 -0700 Subject: [PATCH 117/219] test: a TP leader's early build takes each step's decision once --- test/modular/test_tp_early_spec_decision.py | 94 +++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 test/modular/test_tp_early_spec_decision.py diff --git a/test/modular/test_tp_early_spec_decision.py b/test/modular/test_tp_early_spec_decision.py new file mode 100644 index 0000000000..f3ab1890f9 --- /dev/null +++ b/test/modular/test_tp_early_spec_decision.py @@ -0,0 +1,94 @@ +"""A TP leader's early speculation takes a step's decision exactly once: an +early build that finds a head broadcasts it; one that finds nothing sends no +marker and schedules no yield-away (the top of the next iteration decides). +The TP1 early build still decides on the spot.""" +from types import SimpleNamespace + +from mstar.worker.worker import Worker, _SpecFairness + + +class _Scheduler: + def __init__(self, other_ready=False, yield_batch=None): + self.other_ready = other_ready + self.yield_batch = yield_batch + self.get_next_calls = 0 + + def has_ready_excluding(self, request_state, target): + return self.other_ready + + def get_next_batch(self, request_state, exclude_target=None): + self.get_next_calls += 1 + return self.yield_batch + + +def _leader(head=None, other_ready=False): + w = Worker.__new__(Worker) + w.enable_nvtx = False + w._phase_period = 0 + w.request_state = object() + w.scheduler = _Scheduler(other_ready=other_ready) + w.markers = [] + w.heads = [] + w._can_speculate = lambda batch: True + w._try_speculate_next = lambda pending: head + w._tp_lead_needs_marker = lambda pending, spec: spec is None or spec.tp_seq < 0 + w._broadcast_tp_nospec = lambda pending: w.markers.append(pending.tp_seq) + + def _send(node_batch, speculative=False, spec_from_seq=None): + w.heads.append((speculative, spec_from_seq)) + return 7 + + w.maybe_send_zmq_to_tp_followers = _send + return w + + +def _pending(seq=3): + return SimpleNamespace( + tp_seq=seq, node_name="LLM", graph_walk="decode", + batch=SimpleNamespace(request_to_worker_graph={1: 0}), + node_batch=SimpleNamespace(admit_error=None), + ) + + +def _fair(**kw): + return _SpecFairness(peek=True, hold_s=0.0, max_consecutive=1024, **kw) + + +def test_an_early_build_with_no_head_sends_nothing_and_leaves_the_decision(): + w = _leader(head=None) + spec, target = w._build_speculation(_pending(), 1, _fair(), decide=False) + assert spec is None and target is None + assert w.markers == [] and w.heads == [] and w.scheduler.get_next_calls == 0 + # the full call afterwards takes exactly one decision: the marker + spec, target = w._build_speculation(_pending(), 1, _fair(), decide=True) + assert spec is None and w.markers == [3] and w.heads == [] + assert w.scheduler.get_next_calls == 1 + + +def test_an_early_build_with_a_head_commits_it(): + head = SimpleNamespace(node_batch=object(), tp_seq=-1) + w = _leader(head=head) + spec, target = w._build_speculation(_pending(), 1, _fair(), decide=False) + assert spec is head and spec.tp_seq == 7 and target is None + assert w.heads == [(True, 3)] and w.markers == [] + + +def test_a_fairness_yield_in_the_early_slot_is_deferred_on_a_parallel_node(): + w = _leader(head=SimpleNamespace(node_batch=object(), tp_seq=-1), other_ready=True) + fair = _fair() + spec, target = w._build_speculation(_pending(), 1, fair, decide=False) + # other work is ready and the hold is 0: a yield would be due, but the + # early build is non-committal, so nothing is sent or scheduled + assert spec is None and target is None + assert w.heads == [] and w.markers == [] and w.scheduler.get_next_calls == 0 + # the deciding call yields: marker first, then the yield-away schedule + spec, target = w._build_speculation(_pending(), 1, fair, decide=True) + assert spec is None and target == ("LLM", "decode") + assert w.markers == [3] and w.scheduler.get_next_calls == 1 + + +def test_a_tp1_early_build_decides_on_the_spot(): + w = _leader(head=None) + w._tp_lead_needs_marker = lambda pending, spec: False # not a parallel node + spec, target = w._build_speculation(_pending(), 1, _fair(), decide=True) + assert spec is None and w.markers == [] and w.scheduler.get_next_calls == 1 From 779b9f396f6e87aa292fe8ad9d05346fc45c5888 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:42:54 -0700 Subject: [PATCH 118/219] ci: list the TP early speculation test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 766759e150..1b9862320f 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -59,3 +59,4 @@ test/modular/test_data_worker_wakeup.py test/modular/test_sampler_last_token.py test/modular/test_qwen3_5_device_loopback.py test/modular/test_worker_device_loopback.py +test/modular/test_tp_early_spec_decision.py From e44975b7fdfd3391436d2b651b0e6f57ddbcf69b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 14:44:02 -0700 Subject: [PATCH 119/219] worker: the spec splice runs on the main thread, not the plan thread --- mstar/worker/worker.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index c5ca557ae0..3491e680bd 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2449,9 +2449,10 @@ def _thread_outputs_to_speculative( ): """Splice batch N's outputs into the spec batch's inputs. - Runs on the plan thread as soon as N's outputs are published, so it - must not read a tensor VALUE: it only moves tensor lists and tests - which keys are present. N's kernels may still be running. + Runs on the main thread right after N's future resolves, which is + when the GPU thread has enqueued N, so it must not read a tensor + VALUE: it only moves tensor lists and tests which keys are present. + N's kernels may still be running. """ threaded_continuing: set[int] = set() dropped: set[int] = set() From f5ec4d3b02d81ae911e6b30679df7ec65e4710ab Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 120/219] utils: freeze the set-up heap after a process starts --- mstar/utils/gc_freeze.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 mstar/utils/gc_freeze.py diff --git a/mstar/utils/gc_freeze.py b/mstar/utils/gc_freeze.py new file mode 100644 index 0000000000..062c9ce0b0 --- /dev/null +++ b/mstar/utils/gc_freeze.py @@ -0,0 +1,29 @@ +"""Move a process's set-up heap out of the garbage collector's reach. + +A serving process builds a large, long-lived heap while it starts (configs, +tokenizers, weights' Python wrappers, capture buffers). Every gen-2 +collection afterwards walks all of it: measured at 200 ms per pass in the +API server process, every ~20 s under load, each pass stalling request +admission and token delivery for that long. Collecting once and freezing +moves those objects to the permanent generation, so later collections only +walk what the process created since. Refcounting still frees frozen objects; +only a reference cycle alive at the freeze would be retained for good. +""" +import gc +import logging +import os + +logger = logging.getLogger(__name__) + + +def freeze_after_setup(who: str) -> int: + """``gc.collect()`` then ``gc.freeze()`` for a process that finished its + set-up. Returns the number of objects frozen; 0 and no change when + ``MSTAR_GC_FREEZE=0``.""" + if os.environ.get("MSTAR_GC_FREEZE", "1") != "1": + return 0 + gc.collect() + gc.freeze() + frozen = gc.get_freeze_count() + logger.info("%s: gc.freeze() after setup, %d objects in the permanent generation", who, frozen) + return frozen From f91a2c2b1b754a6c0d63656b742c5cd24c4c5ef4 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 121/219] conductor: freeze the set-up heap once the workers are ready --- mstar/conductor/conductor.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 3569707ef3..42550f4076 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -50,6 +50,7 @@ from mstar.utils.logging_config import quiet_noisy_loggers from mstar.utils.orphan import exit_when_orphaned from mstar.utils.profiler import nvtx_enabled, range_pop, range_push +from mstar.utils.gc_freeze import freeze_after_setup logger = logging.getLogger(__name__) @@ -1537,6 +1538,8 @@ def run(self): from mstar.utils.profiler import range_pop, range_push self._wait_for_workers_ready() + # the set-up heap is complete: keep gen-2 collections off it + freeze_after_setup("Conductor") self.communicator.send( "api_server", From e2b139c6059d9f109ff1d758d7dc465c75b28579 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 122/219] api server: freeze the set-up heap once the workers are ready, one data worker wake per token frame --- mstar/api_server/entrypoint.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 6c9e28e361..aefd843736 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -366,6 +366,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 +581,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 @@ -1148,6 +1157,7 @@ async def _one(request: WarmupRequest, request_id: str) -> None: # The router resolves the loaded model's adapter lazily per request, so models # without an adapter simply return a 404 there and keep working via /generate. from mstar.api_server.openai.router import router as openai_router # noqa: E402 +from mstar.utils.gc_freeze import freeze_after_setup app.include_router(openai_router) From 201e484d64e7b79a1d4bd37e0446e46531b08f22 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 123/219] data worker: a token put without a wake, and a wake for the frame --- mstar/api_server/data_worker.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 16768637de..be0e70c379 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -248,9 +248,12 @@ def new_result_tensors(self, input: ResultTensors): 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 @@ -262,6 +265,11 @@ 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): From 13e42519c75ff077c3908d31df699e705acc010a Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 124/219] docs: MSTAR_GC_FREEZE --- docs/environment_variables.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index d88aa6760f..d9bef285b0 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -335,6 +335,15 @@ Worker scheduling milliseconds lets the arrivals of that window prefill together, which is worth decode throughput at high concurrency and costs up to that much time to first token. + * - ``MSTAR_GC_FREEZE`` + - ``1`` + - Once a process has finished its set-up (the API server and the + conductor when every worker is ready, the worker after warm-up and + graph capture), collect once and ``gc.freeze()`` the heap, so later + gen-2 collections do not walk the long-lived set-up objects. Measured + on the API server process: 200 ms per gen-2 pass every ~20 s under + load before, which stalled request admission and token delivery for + that long. ``0`` leaves the collector alone. * - ``MSTAR_DEVICE_LOOPBACK`` - ``1`` - Whether a decode node may keep its loop-back token on the device: From 294680181198a503761c1a66636ae7577001e342 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 125/219] test: the heap freeze and the per-frame wake --- test/modular/test_gc_freeze.py | 44 ++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 test/modular/test_gc_freeze.py diff --git a/test/modular/test_gc_freeze.py b/test/modular/test_gc_freeze.py new file mode 100644 index 0000000000..bdbb743bc5 --- /dev/null +++ b/test/modular/test_gc_freeze.py @@ -0,0 +1,44 @@ +"""The set-up heap freeze and the data worker's per-frame wake-up.""" +import gc + +from mstar.api_server.data_worker import PreprocessWorker +from mstar.utils.gc_freeze import freeze_after_setup + + +def test_freeze_after_setup_moves_the_heap_and_honours_the_knob(monkeypatch): + gc.unfreeze() + monkeypatch.setenv("MSTAR_GC_FREEZE", "0") + assert freeze_after_setup("test") == 0 and gc.get_freeze_count() == 0 + monkeypatch.setenv("MSTAR_GC_FREEZE", "1") + keep = [[i] for i in range(100)] # live objects the collector tracks + frozen = freeze_after_setup("test") + try: + assert frozen > 0 and gc.get_freeze_count() == frozen + assert keep[0] == [0] # frozen objects are still alive and usable + finally: + gc.unfreeze() + assert gc.get_freeze_count() == 0 + + +def test_a_frame_of_tokens_wakes_the_thread_once(): + facade = PreprocessWorker.__new__(PreprocessWorker) + facade.output_loop_idxs = {"r1": {}, "r2": {}} + facade.per_request_reading_tensors = {"r1": 0, "r2": 0} + + class _Q: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + facade.result_tensor_input_queue = _Q() + signals = [] + facade._signal = lambda: signals.append(1) + facade.new_result_token("r1", 5, None, "text_inputs", "text", wake=False) + facade.new_result_token("r2", 6, None, "text_inputs", "text", wake=False) + assert len(facade.result_tensor_input_queue.items) == 2 and signals == [] + facade.wake() + assert signals == [1] + facade.new_result_token("r1", 7, None, "text_inputs", "text") + assert signals == [1, 1] From 7b2fe58e94ade4cd443b568d757e674e932ba100 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:53:01 -0700 Subject: [PATCH 126/219] ci: list the gc freeze test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 1b9862320f..2427aaa470 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -60,3 +60,4 @@ test/modular/test_sampler_last_token.py test/modular/test_qwen3_5_device_loopback.py test/modular/test_worker_device_loopback.py test/modular/test_tp_early_spec_decision.py +test/modular/test_gc_freeze.py From a8a1e76c65881ae9cfa1039a92eaa9bf9cf7b82e Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:54:09 -0700 Subject: [PATCH 127/219] api server: the gc freeze import belongs with the module imports --- mstar/api_server/entrypoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index aefd843736..6115da607c 100644 --- a/mstar/api_server/entrypoint.py +++ b/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 @@ -1157,7 +1158,6 @@ async def _one(request: WarmupRequest, request_id: str) -> None: # The router resolves the loaded model's adapter lazily per request, so models # without an adapter simply return a 404 there and keep working via /generate. from mstar.api_server.openai.router import router as openai_router # noqa: E402 -from mstar.utils.gc_freeze import freeze_after_setup app.include_router(openai_router) From 79b076e578728d2e7cd5d8bde5f4f4dec536cde0 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 15:54:56 -0700 Subject: [PATCH 128/219] conductor: import block sorted --- mstar/conductor/conductor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index 42550f4076..b8aab3e747 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -33,6 +33,7 @@ from mstar.profile.format import RxInfo, TxInfo from mstar.profile.worker import GraphTimings from mstar.utils.exitcode import describe_exitcode +from mstar.utils.gc_freeze import freeze_after_setup from mstar.utils.ipc_format import ( ConductorMessageType, DrainRequest, @@ -50,7 +51,6 @@ from mstar.utils.logging_config import quiet_noisy_loggers from mstar.utils.orphan import exit_when_orphaned from mstar.utils.profiler import nvtx_enabled, range_pop, range_push -from mstar.utils.gc_freeze import freeze_after_setup logger = logging.getLogger(__name__) From 31dc82cd46d6d044d0e6f71f20fc40af5f56ca8c Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 16:48:41 -0700 Subject: [PATCH 129/219] utils: the device loop-back knob in one place --- mstar/utils/knobs.py | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 mstar/utils/knobs.py diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py new file mode 100644 index 0000000000..ec5ef92f3e --- /dev/null +++ b/mstar/utils/knobs.py @@ -0,0 +1,10 @@ +"""Environment knobs read by more than one package.""" +import os + + +def device_loopback_enabled() -> bool: + """``MSTAR_DEVICE_LOOPBACK`` (default 1): whether a node may keep its + loop-back token on the device (the sampler's slot master) instead of + routing a tensor per request; also whether the sampler persists the + last token at all.""" + return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" From 6422ea8ce0e801e720bbe7b06794968f6eb82f07 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 16:48:41 -0700 Subject: [PATCH 130/219] submodule base: take the loop-back knob from utils --- mstar/model/submodule_base.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 8fd13212b1..542334f581 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Mapping, Sequence @@ -14,6 +13,7 @@ from mstar.communication.tensors import NameToTensorList from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.engine.resources import Resource, SlotLease, SubmoduleStep +from mstar.utils.knobs import device_loopback_enabled # noqa: F401 (re-export) class HostRows(NamedTuple): @@ -259,12 +259,6 @@ def _split_pos_ids( return pos_ids[..., start:end] -def device_loopback_enabled() -> bool: - """``MSTAR_DEVICE_LOOPBACK`` (default 1): whether a node may keep its - loop-back token on the device instead of routing a tensor per request.""" - return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" - - @dataclass class ARNodeInputs(NodeInputs): """ From 0628a7d2685449b120815e908393db1632228167 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 16:48:41 -0700 Subject: [PATCH 131/219] sampler resource: persist the last token only with the loop-back on, time the two scatters --- mstar/engine/resources/sampler/resource.py | 23 ++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index fbed9c20ea..0a4e688e0c 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -1,6 +1,7 @@ from dataclasses import asdict +from time import perf_counter import torch @@ -13,6 +14,8 @@ ) from mstar.engine.resources.sampler.utils import CudaGraphableSampler, Sampler, SamplerBuffers from mstar.engine.resources.step import SlotLease, StepContext +from mstar.utils.knobs import device_loopback_enabled +from mstar.utils.profiler import PHASE_PERIOD, phase_record class SamplerResource(Resource): @@ -73,6 +76,8 @@ def __init__( # (cg slot, padded bs) of the step planned last, for the last-token # read in ``loopback_tokens``; None on an eager step self._step_slot: tuple[int, int] | None = None + # whether the last sampled token is persisted by slot at all + self._keep_last_token = device_loopback_enabled() # rid -> the prompt tokens a cache hit kept out of this step's inputs self._cached_prefix: dict[str, torch.Tensor] = {} @@ -316,8 +321,15 @@ def commit(self, step: SamplerStep, ctx: StepContext): # None on an eager step, which never gathered one if self._cg_buffers is None or self._cg_sampler is None: return + t0 = perf_counter() if PHASE_PERIOD else 0.0 self._cg_buffers.scatter_offset(ctx.slot_lease.slot) - self._cg_buffers.scatter_last_token(ctx.slot_lease.slot) + if PHASE_PERIOD: + phase_record("res.commit.sampler.offset", perf_counter() - t0) + t0 = perf_counter() + if self._keep_last_token: + self._cg_buffers.scatter_last_token(ctx.slot_lease.slot) + if PHASE_PERIOD: + phase_record("res.commit.sampler.last_token", perf_counter() - t0) # Skipped when nothing read the mask this step: there is then nothing to # copy back, and the rows go stale only for requests at penalty 1.0, # which never read them. See `_penalty_live` for why that stays sound. @@ -352,7 +364,10 @@ def sample( # An eager step has no per-step row to scatter from: write the masters # by request, so the next step can still read its token on the device. bufs = self._cg_buffers - if bufs is not None and all(bufs.has_slot(rid) for rid in request_ids): + if ( + self._keep_last_token and bufs is not None + and all(bufs.has_slot(rid) for rid in request_ids) + ): bufs.write_last_tokens(request_ids, tokens) return tokens @@ -360,8 +375,8 @@ def sample( def has_slot_masters(self) -> bool: """Whether the last sampled token of every request is kept on the device by slot (the graph-buffer path): what the device-side - loop-back needs. False on an eager-only node.""" - return self._cg_buffers is not None + loop-back needs. False on an eager-only node or with the knob off.""" + return self._cg_buffers is not None and self._keep_last_token def loopback_tokens(self, request_ids: list[str]) -> torch.Tensor: """The last sampled token of each row of this step, as its input ids, From 253803406ed8ccb67ce49a5c5a3bd75ea634db4b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 16:48:41 -0700 Subject: [PATCH 132/219] qwen3.5: time the loop-back gather --- mstar/model/qwen3_5/submodules.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 581eb4f4a2..a57f51024b 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -9,6 +9,7 @@ import logging import os from collections.abc import Mapping +from time import perf_counter from typing import Any import torch @@ -71,6 +72,7 @@ NodeSubmodule, device_loopback_enabled, ) +from mstar.utils.profiler import PHASE_PERIOD, phase_record logger = logging.getLogger(__name__) @@ -337,7 +339,10 @@ def preprocess( # device loop-back: every row of a decode step is its request's # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] + t0 = perf_counter() if PHASE_PERIOD else 0.0 out["input_ids"] = sampler.loopback_tokens(engine_inputs.request_ids) + if PHASE_PERIOD: + phase_record("engine.preprocess.loopback", perf_counter() - t0) elif inputs[0].input_ids is not None: out["input_ids"] = torch.cat([inp.input_ids for inp in inputs], dim=0) else: From a111f5100fde446183657948da33a29ac1330669 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:08:37 -0700 Subject: [PATCH 133/219] qwen3.5: text prefill graph buckets up to 8192 tokens, with a knob --- mstar/model/qwen3_5/submodules.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index a57f51024b..1d5e83a4c1 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -78,7 +78,17 @@ class LLMSubmodule(ARNodeSubmodule): - PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024, 2048] + # Total tokens of a packed text-prefill step, one captured graph per + # (bucket, batch size). A prompt past the top bucket prefills eagerly, and + # an eager 8k prefill is launch-bound (~1,100 kernel launches, the GPU + # idle half the time: 78 ms of wall for 36 ms of kernels on 0.8B), so the + # ladder runs to 8192. MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS narrows or + # extends it (static buffers are sized by the largest bucket). + PREFILL_TOKEN_BUCKETS = [ + int(x) for x in os.environ.get( + "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", "32,64,128,256,512,1024,2048,4096,8192", + ).split(",") + ] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] # Capture rows and a replay's padding rows address the pool's sink and # hold no slot, so these buckets do not size `gdn_state.max_slots`; that is From 6471b574356685f53881e92c53307edb0ea69fef Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:08:37 -0700 Subject: [PATCH 134/219] docs: MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS --- docs/environment_variables.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index d9bef285b0..09bb9500a3 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -335,6 +335,13 @@ Worker scheduling milliseconds lets the arrivals of that window prefill together, which is worth decode throughput at high concurrency and costs up to that much time to first token. + * - ``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_GC_FREEZE`` - ``1`` - Once a process has finished its set-up (the API server and the From 809af26d3e6f0bf6755792b4b8b2e68c5d3eefff Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:08:37 -0700 Subject: [PATCH 135/219] test: the qwen3.5 bucket ladders come from the environment --- test/modular/test_qwen3_5_bucket_knobs.py | 28 +++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 test/modular/test_qwen3_5_bucket_knobs.py diff --git a/test/modular/test_qwen3_5_bucket_knobs.py b/test/modular/test_qwen3_5_bucket_knobs.py new file mode 100644 index 0000000000..877f36dc73 --- /dev/null +++ b/test/modular/test_qwen3_5_bucket_knobs.py @@ -0,0 +1,28 @@ +"""The Qwen3.5 graph-bucket ladders come from the environment.""" +import importlib + + +def _reload(monkeypatch, name, value): + if value is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, value) + import mstar.model.qwen3_5.submodules as m + return importlib.reload(m) + + +def test_prefill_ladder_default_reaches_8k_and_the_knob_narrows_it(monkeypatch): + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", None) + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS[-1] == 8192 + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS[0] == 32 + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", "64,2048") + assert m.LLMSubmodule.PREFILL_TOKEN_BUCKETS == [64, 2048] + m = _reload(monkeypatch, "MSTAR_QWEN35_PREFILL_TOKEN_BUCKETS", None) + assert len(m.LLMSubmodule.PREFILL_TOKEN_BUCKETS) == 9 + + +def test_decode_ladder_knob(monkeypatch): + m = _reload(monkeypatch, "MSTAR_QWEN35_DECODE_BUCKETS", "1,2,4,8,16,32") + assert m.LLMSubmodule.DECODE_CAPTURE_BATCH_SIZES == [1, 2, 4, 8, 16, 32] + m = _reload(monkeypatch, "MSTAR_QWEN35_DECODE_BUCKETS", None) + assert m.LLMSubmodule.DECODE_CAPTURE_BATCH_SIZES[-1] == 128 From c553dee3caf0083a3f6787b1a6ab1fd73b262aa4 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:08:37 -0700 Subject: [PATCH 136/219] ci: list the bucket knob test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 2427aaa470..0303fef3af 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -61,3 +61,4 @@ test/modular/test_qwen3_5_device_loopback.py test/modular/test_worker_device_loopback.py test/modular/test_tp_early_spec_decision.py test/modular/test_gc_freeze.py +test/modular/test_qwen3_5_bucket_knobs.py From b00e24c8c51079ca69ff39b59db8227dac9cbd81 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:21:57 -0700 Subject: [PATCH 137/219] utils: tokenize a long prompt as parallel chunks, identical ids --- mstar/utils/fast_tokenize.py | 75 ++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 mstar/utils/fast_tokenize.py diff --git a/mstar/utils/fast_tokenize.py b/mstar/utils/fast_tokenize.py new file mode 100644 index 0000000000..fb8c0062e4 --- /dev/null +++ b/mstar/utils/fast_tokenize.py @@ -0,0 +1,75 @@ +"""Tokenize a long prompt as parallel chunks, byte-identical to one encode. + +An 8k-token prompt costs the fast tokenizer ~15 ms single-threaded, a quarter +of the request's time to first token. The Rust backend encodes a batch in +parallel, so a long text is cut into chunks and encoded as one batch; the +cuts are chosen so every pre-token stays whole, which keeps the BPE merges, +and so the ids, identical to encoding the whole text. + +The cut rule assumes a GPT-style pre-tokenizer (the Qwen2 family, GPT-2, +Llama 3): a space starts a pre-token unless it sits inside a run of +whitespace, so a cut is placed only right before a space whose previous +character is not whitespace. Runs of whitespace, newline runs and the +punctuation-plus-newline pre-tokens are never cut, and special-token strings +contain no space. The identity is checked by test/modular/test_fast_tokenize.py +against the real tokenizer over the benchmark prompts. +""" +from __future__ import annotations + +import os + +# Below this many characters a single encode is as fast; above it chunking +# pays. ~4 KB of text is roughly 1k tokens. +DEFAULT_CHUNK_CHARS = int(os.environ.get("MSTAR_TOKENIZE_CHUNK_CHARS", "4096")) +DEFAULT_MAX_CHUNKS = int(os.environ.get("MSTAR_TOKENIZE_MAX_CHUNKS", "8")) + + +def cut_points(text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS) -> list[int]: + """Indices where ``text`` may be cut: each is a space that starts a + whitespace run, at or after the next chunk target. Empty when the text is + short or no safe cut exists.""" + n = len(text) + if chunk_chars <= 0 or n <= chunk_chars: + return [] + want = min(max_chunks, max(2, n // chunk_chars)) + size = n // want + cuts: list[int] = [] + at = 0 + for _ in range(want - 1): + target = at + size + if target >= n: + break + i = text.find(" ", target) + while i > 0 and text[i - 1].isspace(): + i = text.find(" ", i + 1) + if i <= 0 or i >= n - 1: + break + cuts.append(i) + at = i + return cuts + + +def split_at(text: str, cuts: list[int]) -> list[str]: + parts = [] + at = 0 + for c in cuts: + parts.append(text[at:c]) + at = c + parts.append(text[at:]) + return parts + + +def encode_ids(tokenizer, text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS) -> list[int]: + """Token ids of ``text`` with the tokenizer's defaults for special + tokens: a long text goes through the Rust backend as a batch of chunks, + a short one (or a tokenizer without a Rust backend) through ``tokenizer`` + itself.""" + backend = getattr(tokenizer, "backend_tokenizer", None) + cuts = cut_points(text, chunk_chars, max_chunks) if backend is not None else [] + if not cuts: + return list(tokenizer(text).input_ids) + encodings = backend.encode_batch(split_at(text, cuts), add_special_tokens=False) + ids: list[int] = [] + for enc in encodings: + ids.extend(enc.ids) + return ids From d19919fd6679126d8fabd31cbae43a10357a8b1c Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:21:57 -0700 Subject: [PATCH 138/219] qwen3.5: encode the rendered prompt through the chunked path --- mstar/model/qwen3_5/qwen3_5_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index 15844215f5..ddd36b9160 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -60,6 +60,7 @@ Qwen3_5VisionConfig, ) from mstar.model.submodule_base import NodeSubmodule +from mstar.utils.fast_tokenize import encode_ids logger = logging.getLogger(__name__) @@ -421,8 +422,8 @@ def process_prompt( add_generation_prompt=True, enable_thinking=kwargs.get("enable_thinking", True), ) - input_ids = self.tokenizer(text, return_tensors="pt").input_ids[0] - + # a long prompt is encoded as parallel chunks, identical ids + input_ids = torch.tensor(encode_ids(self.tokenizer, text), dtype=torch.long) spans = find_media_spans(input_ids, self._placeholder_specs()) segments = split_around_spans(input_ids, spans) check_plan(prefill_plan(parts), spans, len(segments)) From 723e31fe76ac35698b436222f15f0f79a0ce241c Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:21:57 -0700 Subject: [PATCH 139/219] test: chunked tokenization cuts and identity --- test/modular/test_fast_tokenize.py | 57 ++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 test/modular/test_fast_tokenize.py diff --git a/test/modular/test_fast_tokenize.py b/test/modular/test_fast_tokenize.py new file mode 100644 index 0000000000..df8b35bdab --- /dev/null +++ b/test/modular/test_fast_tokenize.py @@ -0,0 +1,57 @@ +"""Chunked tokenization: cuts only at the start of a whitespace run, chunks +concatenate to the text, and the ids are identical to one encode (checked +against the real Qwen tokenizer when it is in the local cache).""" +import json +import os + +import pytest + +from mstar.utils.fast_tokenize import cut_points, encode_ids, split_at + + +def test_cuts_start_whitespace_runs_and_chunks_rebuild_the_text(): + text = ("word " * 300 + "a b \n\n c. \n d" + " tail " * 300) * 4 + cuts = cut_points(text, chunk_chars=512, max_chunks=8) + assert 1 <= len(cuts) <= 7 + for c in cuts: + assert text[c] == " " and not text[c - 1].isspace() + assert "".join(split_at(text, cuts)) == text + assert cut_points("short text", chunk_chars=4096) == [] + assert cut_points("x" * 10000, chunk_chars=4096) == [] # no space: no safe cut + + +def test_a_short_text_or_a_slow_tokenizer_takes_the_plain_path(): + class _Slow: + def __call__(self, text): + class _R: + input_ids = [len(text)] + return _R() + + assert encode_ids(_Slow(), "x" * 10000) == [10000] + + +def _tokenizer(): + cache = os.environ.get("MSTAR_TEST_HF_HOME") + if not cache: + pytest.skip("set MSTAR_TEST_HF_HOME to a cache holding Qwen/Qwen3.5-0.8B for the identity check") + os.environ.setdefault("HF_HOME", cache) + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B") + + +def test_chunked_ids_are_identical_to_one_encode_on_the_benchmark_prompts(): + tok = _tokenizer() + prompts = os.environ.get("MSTAR_TEST_PROMPTS") + texts = [] + if prompts and os.path.exists(prompts): + d = json.load(open(prompts)) + for key in ("long8k", "long", "short"): + texts += [x if isinstance(x, str) else x.get("prompt", "") for x in d.get(key, [])] + texts += [ + "a b \n\n c. \n d" * 2000, ("<|im_start|>user\nhello world<|im_end|>\n" * 400), + "x,y;z!! \n\nq" * 1500, "tab\tsep arated\r\nlines " * 1200, "数字 123 456 文本 " * 1000, + ] + for t in texts: + whole = tok(t).input_ids + assert encode_ids(tok, t, chunk_chars=2048, max_chunks=8) == whole + assert encode_ids(tok, t, chunk_chars=512, max_chunks=16) == whole From 65f24c93a388de605ae6579fcc4a3893ebcbbd61 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:21:57 -0700 Subject: [PATCH 140/219] ci: list the chunked tokenization test --- test/cpu-core.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 0303fef3af..c2b75c7117 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -62,3 +62,4 @@ test/modular/test_worker_device_loopback.py test/modular/test_tp_early_spec_decision.py test/modular/test_gc_freeze.py test/modular/test_qwen3_5_bucket_knobs.py +test/modular/test_fast_tokenize.py From 9b2415fbab1670f1cfcb51df7353b21e90dc9d7b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:22:42 -0700 Subject: [PATCH 141/219] utils: wrap the chunked tokenizer signatures --- mstar/utils/fast_tokenize.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mstar/utils/fast_tokenize.py b/mstar/utils/fast_tokenize.py index fb8c0062e4..286f61af76 100644 --- a/mstar/utils/fast_tokenize.py +++ b/mstar/utils/fast_tokenize.py @@ -24,7 +24,9 @@ DEFAULT_MAX_CHUNKS = int(os.environ.get("MSTAR_TOKENIZE_MAX_CHUNKS", "8")) -def cut_points(text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS) -> list[int]: +def cut_points( + text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS, +) -> list[int]: """Indices where ``text`` may be cut: each is a space that starts a whitespace run, at or after the next chunk target. Empty when the text is short or no safe cut exists.""" @@ -59,7 +61,10 @@ def split_at(text: str, cuts: list[int]) -> list[str]: return parts -def encode_ids(tokenizer, text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS) -> list[int]: +def encode_ids( + tokenizer, text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, + max_chunks: int = DEFAULT_MAX_CHUNKS, +) -> list[int]: """Token ids of ``text`` with the tokenizer's defaults for special tokens: a long text goes through the Rust backend as a batch of chunks, a short one (or a tokenizer without a Rust backend) through ``tokenizer`` From 148605f5893f9ca5f0ea016cbf2e7d89518aa6a5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:42:14 -0700 Subject: [PATCH 142/219] api: orjson for the SSE events, json fallback --- mstar/api_server/openai/_util.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/openai/_util.py b/mstar/api_server/openai/_util.py index 7ae0f7f240..3481d1b1eb 100644 --- a/mstar/api_server/openai/_util.py +++ b/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: From 53ea6e984794107b82dbc81f8f73f2f42c20ab3f Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:42:14 -0700 Subject: [PATCH 143/219] api: iter_result_chunk_batches, one list per drain --- mstar/api_server/entrypoint.py | 43 ++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 6115da607c..8ab6f3c866 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -743,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 @@ -783,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) @@ -802,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 From ecc6d44d6d0c0eace96c31e7b26a7cadf58ff584 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:42:14 -0700 Subject: [PATCH 144/219] api: chat stream sends a batch of text events as one frame --- mstar/api_server/openai/serving_chat.py | 50 +++++++++++++++---------- 1 file changed, 31 insertions(+), 19 deletions(-) diff --git a/mstar/api_server/openai/serving_chat.py b/mstar/api_server/openai/serving_chat.py index 8e007de3d4..cbb6a24d2d 100644 --- a/mstar/api_server/openai/serving_chat.py +++ b/mstar/api_server/openai/serving_chat.py @@ -97,25 +97,37 @@ def error(message: str, status: int) -> str: yield chunk({"role": "assistant"}) failed = False try: - async for c in api.iter_result_chunks(request_id): - if c.modality == "text": - yield chunk({"content": c.data.decode("utf-8", "replace")}) - elif c.modality == "audio": - # Streaming audio deltas are base64 16-bit PCM at the model rate. - yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) - elif c.modality == "image": - yield chunk({"content": media_io.png_to_data_url(c.data)}) - elif c.modality == "error": - # The request failed after the stream opened (an engine error - # mid generation, a preprocess error); the HTTP status is - # committed, so the failure travels in-band the way the - # non-streaming path's error body does, not as a normal - # ``stop`` a client would take for a complete answer. The error - # is the iterator's last chunk, so the loop ends on its own; - # returning here instead would trip the iterator's abort on a - # request that is already gone. - failed = True - yield error(c.data.decode("utf-8", "replace"), int(c.metadata.get("status", 500))) + # One yield per batch of text chunks: the tokens a request had waiting + # when its handler ran go out as one body frame (one write), which is + # what keeps the delivery loop ahead at high token rates. A batch of + # one is the common case when the loop keeps up. + async for batch in api.iter_result_chunk_batches(request_id): + text_events: list[str] = [] + for c in batch: + if c.modality == "text": + text_events.append(chunk({"content": c.data.decode("utf-8", "replace")})) + continue + if text_events: + yield "".join(text_events) + text_events = [] + if c.modality == "audio": + # Streaming audio deltas are base64 16-bit PCM at the model rate. + yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) + elif c.modality == "image": + yield chunk({"content": media_io.png_to_data_url(c.data)}) + elif c.modality == "error": + # The request failed after the stream opened (an engine error + # mid generation, a preprocess error); the HTTP status is + # committed, so the failure travels in-band the way the + # non-streaming path's error body does, not as a normal + # ``stop`` a client would take for a complete answer. The error + # is the iterator's last chunk, so the loop ends on its own; + # returning here instead would trip the iterator's abort on a + # request that is already gone. + failed = True + yield error(c.data.decode("utf-8", "replace"), int(c.metadata.get("status", 500))) + if text_events: + yield "".join(text_events) except HTTPException as exc: # The delivery timeout raises out of the iterator (which aborts the # request on its way out); report it the same way. From b1a8dc3d6545c3ac2bb673a74a369e2bbc6c0244 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:42:14 -0700 Subject: [PATCH 145/219] test: result chunk batches --- test/modular/test_api_result_delivery.py | 47 ++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/test/modular/test_api_result_delivery.py b/test/modular/test_api_result_delivery.py index d90109f3fa..a0f3786651 100644 --- a/test/modular/test_api_result_delivery.py +++ b/test/modular/test_api_result_delivery.py @@ -303,3 +303,50 @@ def call_soon_threadsafe(self, fn, *args): assert len(loop.calls) == 1 assert all(req.wake.is_set() for req in reqs[:3]) assert not reqs[3].wake.is_set() + + +def test_chunks_waiting_together_come_as_one_batch(): + """Everything the message thread appended since the last wake is drained + as one list, so a streaming handler can send it as one body frame; a + chunk that lands later comes on its own, and the finish ends the stream + without an empty batch.""" + pw = _StubPreprocessWorker() + server = _api_server_stub(pw) + req = _pending_request() + for i in range(3): + req.chunks.append(ResultChunk(request_id="r5", modality="text", data=f"t{i}".encode())) + server.pending_requests["r5"] = req + + async def _collect(): + batches = server.iter_result_chunk_batches("r5") + got = [await batches.__anext__()] + with server.request_lock: + req.chunks.append(ResultChunk(request_id="r5", modality="text", data=b"t3")) + req.notify() + got.append(await batches.__anext__()) + with server.request_lock: + req.event.set() + req.notify() + async for batch in batches: + got.append(batch) + return got + + got = asyncio.run(_collect()) + assert [[c.data for c in b] for b in got] == [[b"t0", b"t1", b"t2"], [b"t3"]] + assert "r5" not in server.pending_requests + + +def test_error_tail_is_its_own_last_batch(): + pw = _StubPreprocessWorker() + server = _api_server_stub(pw) + req = _pending_request() + req.chunks.append(ResultChunk(request_id="r6", modality="text", data=b"partial")) + req.error = "result delivery timed out; response is incomplete" + req.error_status = 500 + req.event.set() + server.pending_requests["r6"] = req + + async def _collect(): + return [[c.modality for c in b] async for b in server.iter_result_chunk_batches("r6")] + + assert asyncio.run(_collect()) == [["text"], ["error"]] From 931ac125a62dbf77dbd73503ffc35ca1fbbbc4b7 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 17:42:14 -0700 Subject: [PATCH 146/219] test: coalesced chat stream frames, sse line --- test/modular/test_openai_router.py | 55 +++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 4 deletions(-) diff --git a/test/modular/test_openai_router.py b/test/modular/test_openai_router.py index 267ddb4a79..2fd333f38f 100644 --- a/test/modular/test_openai_router.py +++ b/test/modular/test_openai_router.py @@ -5,6 +5,7 @@ paths are all exercised without the engine. """ +import asyncio import base64 import json import sys @@ -57,13 +58,16 @@ async def collect_results(self, request_id, raw_request=None): self.last_raw_request = raw_request return self._chunks.get(request_id, []) - async def iter_result_chunks(self, request_id): + async def iter_result_chunk_batches(self, request_id): # Same contract as the real one: a consumer that stops early, or an - # exception, aborts the request; a fully drained stream does not. + # exception, aborts the request; a fully drained stream does not. The + # chunks come the way the server hands them out, in batches (here all + # at once), so the streaming handlers see the coalesced form. finished = False try: - for c in self._chunks.get(request_id, []): - yield c + chunks = self._chunks.get(request_id, []) + if chunks: + yield list(chunks) if self.raise_after is not None: raise self.raise_after finished = True @@ -71,6 +75,15 @@ async def iter_result_chunks(self, request_id): if not finished: self.aborted.append(request_id) + async def iter_result_chunks(self, request_id): + batches = self.iter_result_chunk_batches(request_id) + try: + async for batch in batches: + for c in batch: + yield c + finally: + await batches.aclose() + @pytest.fixture def client_and_stub(monkeypatch): @@ -300,6 +313,40 @@ def test_videos_stream_ndjson(client_and_stub): assert _b64.b64decode(body["data"][0]["b64_json"]) == b"mp4-full" +def test_chat_stream_sends_a_batch_of_tokens_as_one_frame(): + """Tokens that were waiting together leave as one body frame (one write + on the socket), still one SSE event per token for the client; a chunk + of another modality gets a frame of its own.""" + from mstar.api_server.openai.serving_chat import _stream + + class _Batches: + async def iter_result_chunk_batches(self, request_id): + yield [_Chunk("text", b"a"), _Chunk("text", b"b"), _Chunk("text", b"c")] + yield [_Chunk("text", b"d"), _Chunk("audio", b"\x00\x01"), _Chunk("text", b"e")] + + async def _collect(): + return [frame async for frame in _stream(_Batches(), "bagel", "rid", 24000)] + + frames = asyncio.run(_collect()) + # role, abc, d, audio, e, stop, DONE + assert [f.count("data: ") for f in frames] == [1, 3, 1, 1, 1, 1, 1] + events = [json.loads(l[6:]) for f in frames[:-1] for l in f.splitlines() if l.startswith("data: ")] + assert "".join(e["choices"][0]["delta"].get("content", "") for e in events) == "abcde" + assert events[5]["choices"][0]["delta"]["audio"]["data"] == base64.b64encode(b"\x00\x01").decode() + assert events[-1]["choices"][0]["finish_reason"] == "stop" + assert frames[-1] == "data: [DONE]\n\n" + + +def test_sse_event_is_one_json_line(): + from mstar.api_server.openai._util import sse + + obj = {"choices": [{"delta": {"content": "h\u00e9llo \"q\""}}], "n": 1} + line = sse(obj) + assert line.startswith("data: ") and line.endswith("\n\n") + assert "\n" not in line[:-2] + assert json.loads(line[6:-2]) == obj + + def test_chat_stream_reports_a_failed_request_in_band(client_and_stub): """A request that fails after the stream opened ends with an error event, not a ``finish_reason: stop`` that reads as a complete answer.""" From 0b570c5e794650fb2935bea4dde5cf9412c0e893 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 147/219] kv: host lengths iterate as numpy scalars for FlashInfer's max --- mstar/engine/resources/kv/plan.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/plan.py b/mstar/engine/resources/kv/plan.py index c0be0d33a9..3c2549f89d 100644 --- a/mstar/engine/resources/kv/plan.py +++ b/mstar/engine/resources/kv/plan.py @@ -128,7 +128,7 @@ def build_paged_indptrs( paged_kv_indptr=_int32_tensor(kv_indptr), paged_kv_indices=_int32_tensor(all_pages), paged_kv_last_page_len=_int32_tensor(last_page_lens), - kv_lens=_int32_tensor(kv_lens), + kv_lens=_host_lens(kv_lens), ) @@ -136,6 +136,26 @@ def _int32_tensor(values: list[int]) -> torch.Tensor: return torch.from_numpy(np.array(values, dtype=np.int32)) +class HostLens(torch.Tensor): + """A host int32 row whose iteration yields numpy scalars. + + FlashInfer's decode ``plan`` takes the longest row with the Python builtin + ``max(seq_lens)`` and ``.item()`` on the result; iterating a tensor hands + out one 0-d tensor per row, about 2 us each (0.35 ms at 128 rows, on the + plan thread every decode step). Iterating the numpy view costs 0.1 us a + row and the numpy scalar has ``.item()`` too. Everything else FlashInfer + does with the row (``.cpu()``, ``.min()``, ``len``, a ``copy_`` into its + device buffer) is plain tensor behaviour. + """ + + def __iter__(self): + return iter(self.numpy()) + + +def _host_lens(values: list[int]) -> torch.Tensor: + return _int32_tensor(values).as_subclass(HostLens) + + @dataclass class KVPlanOutput: """ From 9cd55c7d364e82123f2c290b99f1bf8a6324a9eb Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 148/219] test: kv plan host lengths --- test/modular/test_kv_plan_host_lens.py | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 test/modular/test_kv_plan_host_lens.py diff --git a/test/modular/test_kv_plan_host_lens.py b/test/modular/test_kv_plan_host_lens.py new file mode 100644 index 0000000000..05fbc8a1bc --- /dev/null +++ b/test/modular/test_kv_plan_host_lens.py @@ -0,0 +1,28 @@ +"""The KV plan's host row of sequence lengths iterates as numpy scalars, so +FlashInfer's ``max(seq_lens).item()`` in its decode plan stays cheap.""" + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from mstar.engine.resources.kv.plan import HostLens, _host_lens # noqa: E402 + + +def test_host_lens_iterates_as_numpy_scalars(): + lens = _host_lens([3, 17, 9]) + assert isinstance(lens, HostLens) and lens.dtype == torch.int32 + items = list(lens) + assert [int(x) for x in items] == [3, 17, 9] + assert all(isinstance(x, np.integer) for x in items) + # what FlashInfer does with it + assert max(lens).item() == 17 and int(lens.min()) == 3 and len(lens) == 3 + assert lens.cpu() is lens + + +def test_host_lens_copies_into_a_plain_buffer(): + lens = _host_lens([5, 6, 7, 8]) + buf = torch.zeros(8, dtype=torch.int32) + buf[: len(lens)].copy_(lens) + assert buf.tolist() == [5, 6, 7, 8, 0, 0, 0, 0] + assert torch.equal(lens, torch.tensor([5, 6, 7, 8], dtype=torch.int32)) From bd5bf7c4c1c2539637e690cec86d480e51345596 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 149/219] runner: admit skips the resources pre_admit covered --- mstar/engine/resources/runner.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index dbc1670c46..e0f606e22c 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -106,6 +106,8 @@ def __init__( # the step whose pre-plan is staged across the pre-planning resources, # as `_step_key` describes it; None when nothing is staged self._staged: tuple | None = None + # (step key, ready) of the last pre_admit whose stage is still up + self._pre_admitted: tuple | None = None def resolve_cached_prefix( self, rid: str, node_name: str, graph_walk: str, @@ -333,9 +335,10 @@ def _step_key(self, step: SubmoduleStep): ), ) - def _drop_stale_preplan(self, step: SubmoduleStep) -> None: + def _drop_stale_preplan(self, step: SubmoduleStep) -> tuple | None: + """Returns the step's key when the stage is for this very step.""" if self._staged is None or step.ctx.is_preplan: - return + return None key = self._step_key(step) if self._staged != key: # Not the normal path: the worker meant to run the staged step @@ -346,6 +349,8 @@ def _drop_stale_preplan(self, step: SubmoduleStep) -> None: self._staged[0], len(self._staged[1]), key[0], len(key[1]), ) self.clear_preplan() + return None + return key @property def staged(self) -> bool: @@ -355,6 +360,7 @@ def staged(self) -> bool: def clear_preplan(self) -> None: """Drop the staged pre-plan on every resource, and the record of it.""" self._staged = None + self._pre_admitted = None for key in self._order: self._resources[key].clear_preplan() @@ -362,10 +368,21 @@ def admit(self, step: SubmoduleStep) -> FullAdmitOutcome: """Reserve capacity for step. A refused admit gives back everything the step reserved. """ - self._drop_stale_preplan(step) + staged_key = self._drop_stale_preplan(step) + # A step whose ``pre_admit`` ran a step ahead, and whose stage is still + # up for it, skips the resources that sweep covered: each would only + # find its own reservation and no-op (0.1 ms of Python per step at + # 128 rows, on the gpu thread). A refusal below still unwinds them. + covered: list[str] = [] ready = True + pre = self._pre_admitted + if pre is not None and staged_key is not None and pre[0] == staged_key: + covered = self._preplan_keys_for(step) + ready = pre[1] admitted: list[str] = [] for key in self._keys_for(step): + if key in covered: + continue if self._nvtx: range_push(f"res.admit.{key}") t0 = perf_counter() if PHASE_PERIOD else 0.0 @@ -381,7 +398,7 @@ def admit(self, step: SubmoduleStep) -> FullAdmitOutcome: "Admit for resource %s failed with error: %s", key, outcome.reason.message ) - for done in reversed([*admitted, key]): + for done in reversed([*covered, *admitted, key]): self._resources[done].rollback_admit( step.get(done), step.ctx, ) @@ -400,6 +417,7 @@ def plan(self, step: SubmoduleStep) -> dict[str, Any]: # the resources promote their staged share below (or plan afresh # if nothing was staged): either way nothing stays staged self._staged = None + self._pre_admitted = None results = step.ctx.plan_results results.clear() for key in self._keys_for(step): @@ -442,9 +460,11 @@ def pre_admit(self, step: SubmoduleStep) -> FullAdmitOutcome: self._resources[done].rollback_admit( step.get(done), step.ctx, ) + self._pre_admitted = None return FullAdmitOutcome(outcome, key) admitted.append(key) ready = ready and outcome.ready + self._pre_admitted = (self._step_key(step), ready) return FULL_ADMIT_OK if ready else FULL_ADMIT_NOT_READY def pre_plan(self, step: SubmoduleStep) -> dict[str, Any]: From e0c3337b80b9aabcab8920364bebec770695feef Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 150/219] engine: phase samples for the pre-plan's declare and admit --- mstar/engine/engine.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 1dc1d31c95..30d82a222c 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1853,12 +1853,15 @@ def _pre_plan_for_batch( batch.step_context.set_padded_rids( cg_runner.step_ids(lease, batch.request_ids) ) + t0 = perf_counter() if PHASE_PERIOD else 0.0 step = submodule_mgmt.submodule.declare_step( graph_walk=batch.step_context.graph_walk, request_ids=batch.step_context.padded_request_ids, inputs=inputs, slot_lease=lease, ) + if PHASE_PERIOD: + phase_record("engine.pre_declare", perf_counter() - t0) if step is None: return False if batch.inputs is not None or ( @@ -1867,14 +1870,17 @@ def _pre_plan_for_batch( batch.step_context.graph_walk, ) ): - # exec drives this step as declared here (its admit ran as - # pre_admit), instead of declaring and admitting again + # exec drives this step as declared here, and its admit skips + # the resources pre_admit covered batch.step = step batch.step_context.is_preplan = True step.set_ctx(batch.step_context) try: + t0 = perf_counter() if PHASE_PERIOD else 0.0 admit_outcome = self._runner.pre_admit(step) + if PHASE_PERIOD: + phase_record("engine.pre_admit", perf_counter() - t0) if not admit_outcome.ok: # exec re-drives the step and reports the failure from there self.reset_pre_plan_for_batch(batch) From 08141f99dbe22580160c1ba02e7bbc2280340a72 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 151/219] worker: postprocess stamps that add up --- mstar/worker/worker.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 3491e680bd..57a75771b9 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -2924,7 +2924,9 @@ def _pp_stage(name: str) -> None: range_pop(synchronize=False) return None - _pp_stage("check_stop") + # the stop check itself is the direct window above; this stage is its + # tail (prefix chains, stream-terminated loops, failed rids) + _pp_stage("check_stop_tail") if self.enable_nvtx: range_pop(synchronize=False) range_push("worker.postprocess.stop_loops", synchronize=False) @@ -3073,8 +3075,7 @@ def _pp_stage(name: str) -> None: if self.enable_nvtx: range_pop(synchronize=False) range_push("worker.postprocess.register_outputs", synchronize=False) - with self._span("worker.postprocess.register_outputs"): - self._register_outputs(route_output) + self._register_outputs(route_output) _pp_stage("register_outputs") return _PostprocessState( From 2c8de516c117cada2e955cf63d244a9dacbcb3c2 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:20:23 -0700 Subject: [PATCH 152/219] test: admit after pre_admit --- test/modular/test_preplan_guard.py | 38 ++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/test/modular/test_preplan_guard.py b/test/modular/test_preplan_guard.py index d60b98e2a4..39532cc5bc 100644 --- a/test/modular/test_preplan_guard.py +++ b/test/modular/test_preplan_guard.py @@ -208,3 +208,41 @@ def test_runner_clear_preplan_reaches_every_resource() -> None: runner.clear_preplan() assert kv.events[-1] == "clear" and attn.events[-1] == "clear" assert runner.plan(_step("a")) == {"kv": "fresh", "attn": "fresh"} + + +class _Counting(_Blind): + """Records admits as well, the way a real resource re-sweeps its rows.""" + + def admit(self, step, ctx): + self.events.append("pre_admit" if ctx.is_preplan else "admit") + return super().admit(step, ctx) + + +def test_admit_of_a_pre_admitted_step_skips_the_resources_that_sweep_covered() -> None: + kv, attn = _Counting(), _Counting(deps=("kv",)) + runner = StepRunner({"kv": kv, "attn": attn}) + staged = _step("a", preplan=True) + assert runner.pre_admit(staged).ok + runner.pre_plan(staged) + + # the same step reaches the GPU thread: nothing to admit again + live = _step("a") + assert runner.admit(live).ok + assert kv.events == ["pre_admit", "pre_plan"] + assert attn.events == ["pre_admit", "pre_plan"] + assert runner.plan(live) == {"kv": "staged", "attn": "staged"} + + # the next step of the same shape was not pre-admitted: a full sweep + assert runner.admit(_step("a")).ok + assert kv.events[-1] == "admit" and attn.events[-1] == "admit" + + +def test_admit_of_a_foreign_step_still_sweeps_everything() -> None: + kv, attn = _Counting(), _Counting(deps=("kv",)) + runner = StepRunner({"kv": kv, "attn": attn}) + staged = _step("a", preplan=True) + assert runner.pre_admit(staged).ok + runner.pre_plan(staged) + assert runner.admit(_step("b", span=17, slot=None)).ok + assert kv.events == ["pre_admit", "pre_plan", "clear", "admit"] + assert attn.events == ["pre_admit", "pre_plan", "clear", "admit"] From ef7eb6694a8541d5e8d6f576575d3d751b036a0f Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:27:24 -0700 Subject: [PATCH 153/219] submodule: host_rows on the batched output --- mstar/model/submodule_base.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 542334f581..add3a19a91 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -85,6 +85,10 @@ class BatchedModelOutput: # of a batch clone: name -> the views, row i for ``row_request_ids[i]``. # The store describes such rows without reading each view. row_views: dict[str, tuple[torch.Tensor, ...]] | None = None + # The stop buffers' host copy, staged by the gpu thread on the execution + # stream behind the step so the step's completion event covers it; the + # main thread reads it after waiting on that event. + host_rows: HostRows | None = None @classmethod def coerce(cls, output: BatchedModelOutput | dict[str, Any]) -> BatchedModelOutput: From f80c7101c4de86f61647f5558319c03fe2a33c2b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:27:24 -0700 Subject: [PATCH 154/219] worker: stage the stop rows' D2H on the gpu thread behind the step --- mstar/worker/worker.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 57a75771b9..270af23cb9 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -234,6 +234,10 @@ class _PostprocessState: signal_idxs: object = None +# pinned host-row buffers in flight at once (see Worker._stage_host_rows) +_HOST_ROWS_RING = 4 + + class _LazyPerRid: """``Worker._rows_to_per_rid`` on demand: the per-rid host dicts are built the first time a consumer reads them. The batched stop check, the row store @@ -313,6 +317,7 @@ def __init__( self.enable_nvtx = nvtx_enabled(enable_nvtx, device) # (node, walk) -> the loop-back signals the node keeps on the device self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} + self._host_rows_slot = 0 # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -1892,6 +1897,7 @@ def _execute_on_gpu_thread( with self._span("worker.gpu_thread.exec"): outputs = engine.exec_and_postprocess(node_batch) if execution_stream is not None: + self._stage_host_rows(outputs) event = torch.Event() event.record(execution_stream) node_batch.completion_event = event @@ -3270,6 +3276,34 @@ def _inline_emit_values( return None, [] return None, [] + def _stage_host_rows(self, outputs: "BatchedModelOutput") -> None: + """Start the device-to-host copy of the step's row-addressed stop + buffers here on the gpu thread, on the execution stream right behind + the clone that made them, so the completion event recorded next + covers it. The main thread's stop check then finds the rows landed + once it has waited on that event: no side stream, no second sync, and + one fewer GIL release on its critical path. + + A small ring of pinned buffers: step N's rows are read by the main + thread while this thread already stages N+1's (the pipeline is one + step deep, the ring leaves room for the early speculation too). + """ + buffers = outputs.check_stop_buffers + if not buffers or outputs.row_request_ids is None: + return + slot = self._host_rows_slot = (self._host_rows_slot + 1) % _HOST_ROWS_RING + host: dict[str, torch.Tensor] = {} + for index, (name, tensor) in enumerate(buffers.items()): + if not (torch.is_tensor(tensor) and tensor.is_cuda): + host[name] = tensor + continue + buf = self._get_pinned_d2h_buffer( + f"host_rows{slot}", tensor.shape, tensor.dtype, index, + ) + buf.copy_(tensor, non_blocking=True) + host[name] = buf + outputs.host_rows = HostRows(tuple(outputs.row_request_ids), host) + def _get_pinned_d2h_buffer( self, purpose: str, @@ -3369,6 +3403,11 @@ def _prematerialize_for_check_stop( list(outputs.row_request_ids) if outputs.row_request_ids is not None else request_ids ) + landed = outputs.host_rows + if landed is not None and completion_event is not None: + # staged by the gpu thread behind the step; the caller has waited + # on the completion event, so the rows are here + return _LazyPerRid(landed.buffers, list(landed.request_ids)), landed if not torch.cuda.is_available() or completion_event is None: if outputs.check_stop_buffers is not None and row_rids is not None: # host tensors already (a CPU device): only the re-keying From 7d9eb7e57e5dd2c13e3934da4d427d1f8ffdf0f5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:27:24 -0700 Subject: [PATCH 155/219] test: staged host rows --- test/modular/test_batched_model_output.py | 59 +++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/test/modular/test_batched_model_output.py b/test/modular/test_batched_model_output.py index 37542448f9..ba35142600 100644 --- a/test/modular/test_batched_model_output.py +++ b/test/modular/test_batched_model_output.py @@ -199,3 +199,62 @@ def test_merging_two_row_addressed_outputs_drops_the_buffers(): )) assert merged.check_stop_buffers is None and merged.row_request_ids is None assert merged.get_check_stop_input()["b"]["new_token"][0].item() == 2 + + +class _StagingWorker(_StubWorker): + """The gpu-thread side of the host rows: stage behind the step, hand + the landed rows to the stop check without a sync of its own.""" + + _host_rows_slot = 0 + _d2h_stream = None + device = torch.device("cpu") + _stage_host_rows = Worker._stage_host_rows + _prematerialize_for_check_stop = Worker._prematerialize_for_check_stop + + +def test_landed_host_rows_skip_the_side_stream_copy(): + """Rows staged by the gpu thread are what the stop check reads: no + side stream, no second copy, row i still request i.""" + rids = ("a", "b", "c") + tokens = torch.tensor([5, 6, 7]) + out = BatchedModelOutput( + check_stop_buffers={"new_token": tokens}, row_request_ids=rids, + ) + worker = _StagingWorker() + worker._stage_host_rows(out) + assert out.host_rows is not None and out.host_rows.request_ids == rids + # a host tensor is passed through as is + assert out.host_rows.buffers["new_token"] is tokens + + per_rid, host_rows = worker._prematerialize_for_check_stop( + out, completion_event=object(), request_ids=list(rids), + ) + assert host_rows is out.host_rows + assert per_rid["b"]["new_token"][0].item() == 6 + assert worker._d2h_stream is None, "no side stream was needed" + + +def test_nothing_to_stage_without_row_buffers(): + out = BatchedModelOutput(per_rid_outputs={"a": {"new_token": [torch.tensor([1])]}}) + worker = _StagingWorker() + worker._stage_host_rows(out) + assert out.host_rows is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_staged_rows_land_behind_the_stream_and_rotate_buffers(): + rids = tuple(f"r{i}" for i in range(6)) + worker = _StagingWorker() + seen = set() + for step in range(3): + tokens = torch.arange(100 * step, 100 * step + 6, device="cuda") + out = BatchedModelOutput( + check_stop_buffers={"new_token": tokens}, row_request_ids=rids, + ) + worker._stage_host_rows(out) + buf = out.host_rows.buffers["new_token"] + assert not buf.is_cuda and buf.is_pinned() + torch.cuda.synchronize() + assert buf.tolist() == list(range(100 * step, 100 * step + 6)) + seen.add(buf.data_ptr()) + assert len(seen) == 3, "consecutive steps use different pinned buffers" From 56dcfa8ee5f20fff1bd13f6df3e2afa9d8c8aa41 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:50 -0700 Subject: [PATCH 156/219] submodule: rows_only on the batched output, hook docs --- mstar/model/submodule_base.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index add3a19a91..1600a026e1 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -89,6 +89,9 @@ class BatchedModelOutput: # stream behind the step so the step's completion event covers it; the # main thread reads it after waiting on that event. host_rows: HostRows | None = None + # No per-rid entries were built: every output signal of the step stays on + # the device (see ``ExecutingBatch.rows_only``). Readers use the rows. + rows_only: bool = False @classmethod def coerce(cls, output: BatchedModelOutput | dict[str, Any]) -> BatchedModelOutput: @@ -703,8 +706,11 @@ def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: routes them with no tensor: the next step reads the value in ``preprocess`` and every fallback batch arrives with the signal empty. Only sound when ``prepare_inputs`` accepts the empty signal on - this walk. Empty by default; ``MSTAR_DEVICE_LOOPBACK=0`` keeps every - node on the tensor path.""" + this walk. On a walk whose output signals all stay on the device the + engine builds no per-request output dicts and does not call + ``postprocess``: the stop check, the inline emit and the next step's + speculation read the step's rows instead. Empty by default; + ``MSTAR_DEVICE_LOOPBACK=0`` keeps every node on the tensor path.""" del graph_walk return frozenset() From 6f2669e4e6c78cd155422fe6c271e677704ea9fc Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:50 -0700 Subject: [PATCH 157/219] engine: rows-only collect on a device loop-back step, no tail --- mstar/engine/engine.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 30d82a222c..c7b88b573d 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -295,6 +295,12 @@ class ExecutingBatch: # the resource that ran out, so an eviction can be scoped to it failed_resource: str | None = None + # Every output signal of this step stays on the device (the node's + # ``device_loopback_signals``): collect builds no per-request output + # dicts and the per-request ``postprocess`` is skipped; the stop check, + # the inline emit and the speculation splice read the rows instead. + rows_only: bool = False + # This step's outputs, published as soon as the forward has been # submitted — the tensors exist then, even though their values land later. outputs: BatchedModelOutput = field(default_factory=BatchedModelOutput) @@ -973,6 +979,7 @@ def _exec_single( submodule_mgmt, lease, raw, inputs, req_info, request_ids=batch.request_ids, step_request_ids=batch.step_context.padded_request_ids, + rows_only=batch.rows_only, ) finally: if PHASE_PERIOD: @@ -1391,7 +1398,8 @@ def exec_and_postprocess( outputs = self.exec(batch) # A failed admit means no forward ran and every rid's outputs are empty, # so the tail has nothing to consume; the worker re-drives the step. - if batch.admit_error is None: + # A rows-only step has no per-request outputs for a tail to work on. + if batch.admit_error is None and not batch.rows_only: self.postprocess_batch(batch, outputs) return outputs @@ -1536,10 +1544,16 @@ def _collect_outputs( req_info: Mapping[str, CurrentForwardPassInfo], request_ids: list[str], step_request_ids: tuple[str, ...], + rows_only: bool = False, ) -> BatchedModelOutput: """Per-rid outputs for the real requests: drop the padding rows and map a captured graph's keys back to real ids. + ``rows_only`` (see ``ExecutingBatch.rows_only``) keeps the row views + and the stop buffers but builds no per-rid dicts: nothing downstream + reads them on such a step, and at 128 rows they were a few hundred + allocations on the gpu thread. + A captured forward emits its per-rid entries under the slot's padding ids (those were the batch at capture time), so entry ``i`` belongs to ``request_ids[i]`` on either path. @@ -1553,6 +1567,7 @@ def _collect_outputs( row_clones, row_views = self._merge_per_rid( outputs, raw_outputs, request_ids, out_ids, submodule, req_info, + rows_only=rows_only, ) self._merge_unpacked( outputs, raw_outputs, request_ids, submodule, @@ -1560,6 +1575,7 @@ def _collect_outputs( ) return BatchedModelOutput( per_rid_outputs=outputs, + rows_only=rows_only and not outputs, # the stop buffer that is a row output shares its clone: one copy # of the sampled tokens per step, not two check_stop_buffers=raw_outputs.clone_check_stop_buffers( @@ -1579,6 +1595,7 @@ def _merge_per_rid( out_ids: list[str], submodule: NodeSubmodule, req_info: Mapping[str, CurrentForwardPassInfo], + rows_only: bool = False, ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[torch.Tensor, ...]] | None]: """Fold the forward's per-rid entries into ``outputs``. @@ -1605,7 +1622,7 @@ def _merge_per_rid( for name, tensor in row_clones.items() if isinstance(tensor, torch.Tensor) and tensor.dim() } - if row_views: + if row_views and not rows_only: for i, rid in enumerate(request_ids): merged = outputs.setdefault(rid, {}) for name, views in row_views.items(): From 6798a40a414c16b77f503ecead6339a0ba6a49e5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:51 -0700 Subject: [PATCH 158/219] worker: rows-only steps, splice the loop-back signal as present --- mstar/worker/worker.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 270af23cb9..a1ad88b17e 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -1843,6 +1843,12 @@ def _execute_on_gpu_thread( from mstar.utils.profiler import range_pop, range_push engine = self.engine_manager.get_engine(batch.node_name) + # Every output of this step stays on the device: the engine keeps the + # rows and skips the per-request dicts and tail (cached per node and + # walk, a dict read here). + node_batch.rows_only = self._all_on_device( + engine, batch.node_name, batch.graph_walk, batch.output_signals, + ) if logger.isEnabledFor(logging.DEBUG): # test/waypoint/serve_rollout.py parses this line (wire ids, logged # before prepare_inputs runs) to recover the DiT schedule. @@ -2462,9 +2468,23 @@ def _thread_outputs_to_speculative( """ threaded_continuing: set[int] = set() dropped: set[int] = set() + # A rows-only step kept every output on the device: a rid that ran is + # present, and its loop-back signal arrives empty, the way every + # fallback batch of such a walk does (``device_loopback_signals``). + rows_only = getattr(outputs_N, "rows_only", False) + ran = frozenset(outputs_N.row_request_ids or ()) if rows_only else None for rid in list(speculation.node_batch.request_ids): if rid not in speculation.continuing_rids: continue # fresh rid — inputs already gathered. + if ran is not None: + if rid in ran: + inputs = speculation.node_batch.per_request_input_tensors[rid] + for input_name, _ in speculation.consumed_edges: + inputs[input_name] = [] + threaded_continuing.add(rid) + else: + dropped.add(rid) + continue rid_outputs = outputs_N.get(rid, {}) ok = True for input_name, _ in speculation.consumed_edges: From a98f888f7b1f9ee76ccf387f7f0230b4cf2520c9 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:51 -0700 Subject: [PATCH 159/219] test: rows-only collect --- .../modular/test_collect_outputs_fast_path.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/test/modular/test_collect_outputs_fast_path.py b/test/modular/test_collect_outputs_fast_path.py index db3a049aa6..7db8fd6a88 100644 --- a/test/modular/test_collect_outputs_fast_path.py +++ b/test/modular/test_collect_outputs_fast_path.py @@ -118,3 +118,22 @@ def test_stop_buffer_shares_the_row_clone(): ) stop2 = other.clone_check_stop_buffers(reuse=other.clone_row_outputs(2)) assert torch.equal(stop2["new_token"], torch.ones(2, dtype=torch.int64)) + + +def test_rows_only_keeps_the_views_and_builds_no_per_rid_dicts(): + """A step whose outputs all stay on the device: the row views and the + stop buffer are what its readers use, so no per-rid dict is built.""" + rids = [1, 2, 3, 4] + raw = _raw(bs=4) + eng = Engine.__new__(Engine) + outputs = {} + row_clones, row_views = eng._merge_per_rid( + outputs, raw, rids, rids, _Sub(), {}, rows_only=True, + ) + assert outputs == {} + assert set(row_views) == {"new_token", "aux"} and len(row_views["new_token"]) == 4 + assert row_clones["new_token"].tolist() == [100, 101, 102, 103] + # the ordinary path still fills them + outputs = {} + eng._merge_per_rid(outputs, _raw(bs=4), rids, rids, _Sub(), {}) + assert set(outputs) == set(rids) From c5ded392c5d7e4a77fab2a9e0f4fc9b2f277013c Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:51 -0700 Subject: [PATCH 160/219] test: stub carries rows_only --- test/modular/test_admit_failure_handling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_admit_failure_handling.py b/test/modular/test_admit_failure_handling.py index 97fcdedd34..8eb987489e 100644 --- a/test/modular/test_admit_failure_handling.py +++ b/test/modular/test_admit_failure_handling.py @@ -74,7 +74,7 @@ def postprocess_batch(self, batch, outputs): def _batch(): - return SimpleNamespace(request_ids=["r0", "r1"], admit_error=None) + return SimpleNamespace(request_ids=["r0", "r1"], admit_error=None, rows_only=False) @pytest.mark.parametrize( From cdba6fd4d9d07bd2b259c15f13dfee61293624f3 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:31:51 -0700 Subject: [PATCH 161/219] test: rows-only splice --- test/modular/test_multi_stream_final_flag.py | 24 ++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/test/modular/test_multi_stream_final_flag.py b/test/modular/test_multi_stream_final_flag.py index f8bd466bf3..099e33785b 100644 --- a/test/modular/test_multi_stream_final_flag.py +++ b/test/modular/test_multi_stream_final_flag.py @@ -160,3 +160,27 @@ def test_a_dropped_speculative_rid_neither_flushes_nor_reports_done(): assert speculation.dropped == {"r"} assert (batch.final_stream_rids, batch.stream_partition_done_rids) == (set(), set()) assert _step(worker, "vocoder", {"r": {"pitch"}}) == ({"r"}, {"r"}) + + +def test_a_rows_only_step_threads_its_loop_back_signal_as_present(): + """With every output of the step on the device, a continuing rid that + ran is spliced with an empty loop-back signal (what the walk's fallback + batches carry too); one that did not run is dropped.""" + worker = _worker() + batch = Worker._make_executing_batch( + worker, node_name="vocoder", graph_walk="chunk", request_ids=["r", "s"], + per_request_input_tensors={"r": {}, "s": {}}, + per_request_info={"r": None, "s": None}, + ) + speculation = SimpleNamespace( + node_batch=batch, continuing_rids={"r", "s"}, consumed_edges=[("loop", None)], + scheduled_batch=SimpleNamespace(request_to_worker_graph={"r": 0, "s": 0}), + consumed_streaming_edges={}, spec_id=0, + ) + outputs_N = SimpleNamespace(rows_only=True, row_request_ids=("r",)) + + Worker._thread_outputs_to_speculative(worker, speculation, outputs_N) + + assert speculation.continuing_rids == {"r"} + assert batch.per_request_input_tensors["r"] == {"loop": []} + assert batch.request_ids == ["r"] and "s" not in batch.per_request_input_tensors From bd75d9d80294c4d5c09403e0c73c81a8ce205d25 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:34:55 -0700 Subject: [PATCH 162/219] test: stub carries output_signals --- test/modular/test_preplan_error_paths.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_preplan_error_paths.py b/test/modular/test_preplan_error_paths.py index db9e9a1467..bf870038bd 100644 --- a/test/modular/test_preplan_error_paths.py +++ b/test/modular/test_preplan_error_paths.py @@ -131,7 +131,7 @@ def test_gpu_thread_drops_the_stage_when_prepare_inputs_raises() -> None: request_ids=("a",), release_waiters=lambda: released.append(True), per_request_info={"a": SimpleNamespace(request_id="a")}, ) - batch = SimpleNamespace(node_name="dit", graph_walk="decode") + batch = SimpleNamespace(node_name="dit", graph_walk="decode", output_signals=()) with pytest.raises(RuntimeError, match="bad inputs"): Worker._execute_on_gpu_thread(w, batch, node_batch, _done_future()) assert resets == [node_batch] From 633a284e73ef6c11435cafa0cd20ba826c88949d Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:38:35 -0700 Subject: [PATCH 163/219] engine: no clone on a rows-only step --- mstar/engine/engine.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index c7b88b573d..21db132473 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1559,6 +1559,24 @@ def _collect_outputs( ``request_ids[i]`` on either path. """ submodule = submodule_mgmt.submodule + if ( + rows_only + and raw_outputs.check_stop_buffers is not None + and not raw_outputs.per_rid_outputs + and not raw_outputs.packed_outputs + and type(submodule).filter_batched_output + is NodeSubmodule.filter_batched_output + ): + # No clone either: the worker takes the host copy on this thread, + # on this stream, right behind the step (``_stage_host_rows``), so + # a later replay cannot overwrite the rows first, and nothing else + # reads them on the device. Padded rows past the real ones ride + # along and are never read. + return BatchedModelOutput( + check_stop_buffers=raw_outputs.check_stop_buffers, + row_request_ids=tuple(request_ids), + rows_only=True, + ) out_ids = ( step_request_ids if lease is None else submodule_mgmt.cuda_graph_runner.slot_for(lease).dummy_rids From 627fc0e5e5ca8cfe2773aace9f4d1f3626806d2b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:38:35 -0700 Subject: [PATCH 164/219] sampler: memoized per-step views for the commit scatters --- mstar/engine/resources/sampler/utils.py | 39 +++++++++++++++++++------ 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index b1c9d50d55..7e3b5bff19 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1105,6 +1105,7 @@ class Buffer: master: torch.Tensor default: float dtype: torch.dtype + _slot_views: dict = field(default_factory=dict) @classmethod def allocate( @@ -1133,8 +1134,18 @@ def grow_master(self, new_capacity: int) -> None: def slot_view(self, cg_slot: int, bs: int) -> torch.Tensor: """This slot's per-step row (clamped: a single-buffered buffer — the - RNG offset, gathered inline and used serialized — ignores ``cg_slot``).""" - return self.buf[cg_slot if self.buf.shape[0] > 1 else 0, :bs] + RNG offset, gathered inline and used serialized — ignores ``cg_slot``). + + Memoized per (slot, rows): ``buf`` is never reallocated, and the view + is taken on the gpu thread every step, where each tensor call is a + GIL release the main thread's postprocess can win.""" + key = (cg_slot, bs) + view = self._slot_views.get(key) + if view is None: + view = self._slot_views[key] = self.buf[ + cg_slot if self.buf.shape[0] > 1 else 0, :bs + ] + return view def gather(self, idx_view: torch.Tensor, padded_bs: int, cg_slot: int) -> None: torch.index_select(self.master, 0, idx_view, out=self.slot_view(cg_slot, padded_bs)) @@ -1146,7 +1157,7 @@ def scatter(self, idx_view: torch.Tensor, real_bs: int, cg_slot: int) -> None: graph (the RNG offset). REAL rows only: padding rows all gather from slot 0 and get advanced too, so scattering them would clobber slot 0. """ - self.master.index_copy_(0, idx_view[:real_bs], self.slot_view(cg_slot, real_bs)) + self.master.index_copy_(0, idx_view, self.slot_view(cg_slot, real_bs)) @dataclass @@ -1314,6 +1325,8 @@ class SamplerBuffers: # Real (unpadded) batch size of the last gather per cg slot — the # scatter-back writes only these rows (padding rows all map to slot 0). _last_real_bs: list[int] = field(default_factory=list, repr=False) + # memoized per (slot, real rows): see ``_scatter_idx`` + _scatter_idx_views: dict = field(default_factory=dict, repr=False) # cg slot -> the (rids, padded_bs) its pinned index row currently holds, so # gather_dynamic can re-issue the H2D without redoing the O(bs) CPU fill _staged_rids: dict[int, tuple[tuple[str, ...], int]] = field( @@ -1679,13 +1692,22 @@ def sampler_for(self, padded_bs: int, cg_slot: int) -> "CudaGraphableSampler": sampler.applied_penalty_in_graph = False return sampler + def _scatter_idx(self, cg_slot: int) -> tuple[torch.Tensor, int]: + """The slot's device index row narrowed to the real rows, memoized per + (slot, rows) like ``Buffer.slot_view`` (the index buffer is static).""" + real_bs = self._last_real_bs[cg_slot] + key = (cg_slot, real_bs) + idx = self._scatter_idx_views.get(key) + if idx is None: + idx = self._scatter_idx_views[key] = self._slot_idx_gpu[cg_slot, :real_bs] + return idx, real_bs + def scatter_offset(self, cg_slot: int = 0) -> None: """Persist the (in-graph advanced) per-step offsets back to their slot masters. Call once AFTER the graph replay for the gather on ``cg_slot``; GPU-only, real rows only (padding rows all map to slot 0).""" - self.offset.scatter( - self._slot_idx_gpu[cg_slot], self._last_real_bs[cg_slot], cg_slot - ) + idx, real_bs = self._scatter_idx(cg_slot) + self.offset.scatter(idx, real_bs, cg_slot) # ------------------------------------------------------------------ # Last sampled token: written by the sample, read by the next step @@ -1694,9 +1716,8 @@ def scatter_last_token(self, cg_slot: int = 0) -> None: """Persist the tokens the captured sample wrote into ``cg_slot``'s per-step row to their slot masters. Same contract as ``scatter_offset``: after the replay, real rows only.""" - self.last_token.scatter( - self._slot_idx_gpu[cg_slot], self._last_real_bs[cg_slot], cg_slot - ) + idx, real_bs = self._scatter_idx(cg_slot) + self.last_token.scatter(idx, real_bs, cg_slot) def gather_last_tokens(self, cg_slot: int, padded_bs: int) -> torch.Tensor: """``[padded_bs]`` last tokens of the batch staged on ``cg_slot``, in From 727d3bd7d59fee61c665a9e7ae29ae789d2600d9 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:38:35 -0700 Subject: [PATCH 165/219] test: rows-only collect takes no clone --- test/modular/test_collect_outputs_fast_path.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/modular/test_collect_outputs_fast_path.py b/test/modular/test_collect_outputs_fast_path.py index 7db8fd6a88..afb9bc8328 100644 --- a/test/modular/test_collect_outputs_fast_path.py +++ b/test/modular/test_collect_outputs_fast_path.py @@ -137,3 +137,20 @@ def test_rows_only_keeps_the_views_and_builds_no_per_rid_dicts(): outputs = {} eng._merge_per_rid(outputs, _raw(bs=4), rids, rids, _Sub(), {}) assert set(outputs) == set(rids) + + +def test_rows_only_collect_takes_no_clone(): + """On a rows-only step the gpu thread copies the rows to the host right + behind the step, so the graph's own output buffer is handed over as is.""" + from types import SimpleNamespace + + rids = [1, 2, 3, 4] + raw = _raw(bs=4) + eng = Engine.__new__(Engine) + out = eng._collect_outputs( + SimpleNamespace(submodule=_Sub(), cuda_graph_runner=None), None, raw, + [], {}, request_ids=rids, step_request_ids=tuple(rids), rows_only=True, + ) + assert out.rows_only and out.per_rid_outputs == {} and out.row_views is None + assert out.check_stop_buffers["new_token"] is raw.check_stop_buffers["new_token"] + assert out.row_request_ids == (1, 2, 3, 4) From 0014739d0e44b937425884c26b27c032fa6df5a0 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:43:48 -0700 Subject: [PATCH 166/219] worker: skip the stream polls of a speculation build when no request owns a buffer --- mstar/worker/worker.py | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index a1ad88b17e..1d57d67fb2 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -291,6 +291,11 @@ class Worker: computation via engines. """ + # Rids that own a stream buffer (a streaming model's chunked inputs). + # Empty for a text model, which lets a speculation build skip its three + # per-row stream polls; None (a bare instance) means unknown, so poll. + _stream_buffer_rids: set[int] | None = None + def __init__( self, worker_id: str, @@ -318,6 +323,7 @@ def __init__( # (node, walk) -> the loop-back signals the node keeps on the device self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} self._host_rows_slot = 0 + self._stream_buffer_rids = set() # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -725,6 +731,8 @@ def _add_new_request(self, body: NewRequest) -> None: req_info.stream_buffers[conn.edge_name] = sbuf consumer = self._consumer_node_cache.get(conn.edge_name, "") req_info.stream_buffers_by_consumer.setdefault(consumer, {})[conn.edge_name] = sbuf + if self._stream_buffer_rids is not None: + self._stream_buffer_rids.add(request_id) if conn.to_partition == body.request_info.partition_name: lead = body.request_info.stream_lead_items.get(conn.edge_name) if lead: @@ -820,6 +828,8 @@ def _remove_request(self, body: RemoveRequest) -> None: self._pending_drains.discard(body.request_id) self._reads_done_sent.discard(body.request_id) self.engine_manager.remove_request(request_id) + if self._stream_buffer_rids: + self._stream_buffer_rids.discard(request_id) self.request_state.remove_request(request_id) self.tensor_manager.force_cleanup_request(request_id) self.profile_info.pop_request(request_id) @@ -1198,6 +1208,12 @@ def _mark_stream_ingested( if sbuf is not None: sbuf.ingested_chunk = streaming_edge.chunk + def _streams_possible(self) -> bool: + """False when no live request owns a stream buffer: the per-row stream + polls of a speculation build then have nothing to find.""" + rids = getattr(self, "_stream_buffer_rids", None) + return rids is None or bool(rids) + def _stream_chunks_for( self, request_id: int, node_name: str, inputs: NameToTensorList ) -> dict[str, StreamChunkInfo] | None: @@ -1365,11 +1381,12 @@ def _build_input_metadata_and_fwd_info( ) -> tuple[dict[str, InputMetadata], dict[str, CurrentForwardPassInfo]]: meta = {} fwd_info = {} + streams = self._streams_possible() for rid, inputs in per_request_inputs.items(): req_info = self.request_state.get_fwd_info(rid, batch_partition) fwd_info[rid] = req_info stream_chunks = {} - if chunks := self._stream_chunks_for(rid, node_name, inputs): + if streams and (chunks := self._stream_chunks_for(rid, node_name, inputs)): stream_chunks = chunks meta[rid] = InputMetadata( stream_chunks=stream_chunks, @@ -2186,7 +2203,7 @@ def _assemble_speculation( per_request_input_metadata = { rid: InputMetadata(stream_chunks=chunks) for rid in request_ids if (chunks := self._stream_chunks_for(rid, spec_node, per_request_inputs[rid])) - } + } if self._streams_possible() else {} spec_batch = ScheduledBatch( node_name=spec_node, graph_walk=pending.graph_walk, @@ -2361,10 +2378,13 @@ def _try_speculate_next( # Polling the StreamBuffers stays on this side: they hold real tensors. polled: list[tuple[int, StreamingEdge]] = [] per_rid_counts: list[int] = [] - for r in candidates: - edges = self._poll_stream_buffers_for_speculation(r, spec_node_name) - polled.extend((r, e) for e in edges) - per_rid_counts.append(len(edges)) + if self._streams_possible(): + for r in candidates: + edges = self._poll_stream_buffers_for_speculation(r, spec_node_name) + polled.extend((r, e) for e in edges) + per_rid_counts.append(len(edges)) + else: + per_rid_counts = [0] * len(candidates) prep = self._graph_runtime.prep_spec_rids(SpeculationPrepInput( spec_node_name=spec_node_name, From 30eee5a51cff7fb58a8ced5557b18838ca19bce5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:43:48 -0700 Subject: [PATCH 167/219] test: stream polls skipped only without buffers --- test/modular/test_multi_stream_final_flag.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/modular/test_multi_stream_final_flag.py b/test/modular/test_multi_stream_final_flag.py index 099e33785b..ea712c308c 100644 --- a/test/modular/test_multi_stream_final_flag.py +++ b/test/modular/test_multi_stream_final_flag.py @@ -184,3 +184,14 @@ def test_a_rows_only_step_threads_its_loop_back_signal_as_present(): assert speculation.continuing_rids == {"r"} assert batch.per_request_input_tensors["r"] == {"loop": []} assert batch.request_ids == ["r"] and "s" not in batch.per_request_input_tensors + + +def test_stream_polls_are_skipped_only_when_no_request_owns_a_buffer(): + worker = _worker() + assert Worker._streams_possible(worker), "unknown (a bare instance): poll" + worker._stream_buffer_rids = None + assert Worker._streams_possible(worker) + worker._stream_buffer_rids = set() + assert not Worker._streams_possible(worker) + worker._stream_buffer_rids.add(7) + assert Worker._streams_possible(worker) From a717532f67bcceea27a1bb0c74b7cecba4e28113 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:46:26 -0700 Subject: [PATCH 168/219] submodule: uniform_row_inputs hook --- mstar/model/submodule_base.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 1600a026e1..77ea8862b2 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -691,6 +691,16 @@ def filter_batched_output( a fixed key set for graph compat, so the filtering happens here.""" return outputs + def uniform_row_inputs(self, graph_walk: str) -> "NodeInputs | None": + """The inputs every row of this walk gets, when ``prepare_inputs`` + would return an equal object for each regardless of the request, its + routed signals and its metadata (a decode walk that reads its token + off the device). The engine then prepares the batch with one call and + shares the object across the step's rows, so it must not be mutated + per row. None (the default) prepares row by row.""" + del graph_walk + return None + def step_is_input_free(self, graph_walk: str) -> bool: """Whether ``declare_step`` for this walk reads nothing from the inputs but their spans, and every real row has the span of the template row. From 4243dadd3c86a901313ce33f62e894dfa1c8f66a Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:46:26 -0700 Subject: [PATCH 169/219] qwen3_5: the loop-back decode row is uniform --- mstar/model/qwen3_5/submodules.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 1d5e83a4c1..b1a1f0d77d 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -296,6 +296,12 @@ def prepare_inputs( input_ids = inputs["text_inputs"][0] return ARNodeInputs(input_seq_len=input_ids.shape[0], input_ids=input_ids) + def uniform_row_inputs(self, graph_walk: str) -> ARNodeInputs | None: + # the loop-back decode row above, built once for the step's rows + if graph_walk == "decode" and self._device_loopback(): + return ARNodeInputs(input_seq_len=1) + return None + def _device_loopback(self) -> bool: return ( device_loopback_enabled() From 912b2e0a0d703febe832a5d2a13c71fc6e69679f Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:46:26 -0700 Subject: [PATCH 170/219] engine: prepare a uniform walk with one object --- mstar/engine/engine.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 21db132473..54c52887b0 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -760,6 +760,18 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: # read leniently, as the probe itself was only reached per prepared row keyed = getattr(self, "_keyed_walks", None) probe_prefix = bool(keyed) and walk in keyed.get(batch.node_name, ()) + if not probe_prefix: + uniform = submodule.uniform_row_inputs(walk) + if uniform is not None: + # one object for all the step's rows: the per-row call would + # return an equal one each time (128 calls at a full decode + # batch, on the gpu thread before the launch) + node_inputs = [uniform] * len(batch.request_ids) + batch.register_prepare_batch(node_inputs) + batch.running_batched = submodule.can_batch( + batch=batch, model_inputs=node_inputs + ) + return for rid in batch.request_ids: try: req_inputs = prepare( From 9191dcc8ad84ffcb0dd3ed75fb56016b83b98a71 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:46:26 -0700 Subject: [PATCH 171/219] test: uniform loop-back row --- test/modular/test_qwen3_5_device_loopback.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py index dedab1d856..d98b282eb9 100644 --- a/test/modular/test_qwen3_5_device_loopback.py +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -78,3 +78,16 @@ def _position_ids_3d(inputs): assert out["input_ids"].tolist() == [3, 4, 0, 0] assert sampler.asked == ["r1", "r2", "pad1", "pad2"] assert "input_embeds" not in out and seen["n"] == 4 + + +def test_the_loop_back_decode_row_is_uniform(monkeypatch): + from mstar.model.submodule_base import ARNodeInputs + + sub = _sub(_Sampler(masters=True)) + row = LLMSubmodule.uniform_row_inputs(sub, "decode") + assert isinstance(row, ARNodeInputs) and row.input_seq_len == 1 + assert row.input_ids is None and row.input_embeds is None + assert LLMSubmodule.uniform_row_inputs(sub, "prefill_text") is None + assert LLMSubmodule.uniform_row_inputs(_sub(_Sampler(masters=False)), "decode") is None + monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") + assert LLMSubmodule.uniform_row_inputs(sub, "decode") is None From 1e8efa0fdb109d9a22f990a39430a5dba8cef952 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 18:46:26 -0700 Subject: [PATCH 172/219] test: engine uniform inputs --- test/modular/test_engine_uniform_inputs.py | 74 ++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 test/modular/test_engine_uniform_inputs.py diff --git a/test/modular/test_engine_uniform_inputs.py b/test/modular/test_engine_uniform_inputs.py new file mode 100644 index 0000000000..7adeca9247 --- /dev/null +++ b/test/modular/test_engine_uniform_inputs.py @@ -0,0 +1,74 @@ +"""Engine.prepare_inputs: a walk whose rows all get the same inputs is +prepared with one call, the object shared by the rows; any other walk still +prepares row by row.""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.engine import Engine # noqa: E402 + + +class _Batch: + def __init__(self, rids, walk="decode"): + self.node_name = "llm" + self.request_ids = list(rids) + self.step_context = SimpleNamespace(graph_walk=walk) + self.per_request_info_wrapped = {r: SimpleNamespace() for r in rids} + self.per_request_input_tensors = {} + self.per_request_input_metadata = {} + self.final_stream_rids = set() + self.skipped_rids = set() + self.failed_requests = {} + self.inputs = None + self.running_batched = None + + def register_prepare_batch(self, inputs): + self.inputs = inputs + + def drop_rids(self, rids): + self.request_ids = [r for r in self.request_ids if r not in rids] + + def register_failure(self, rid, error): + self.failed_requests[rid] = str(error) + + +class _Sub: + def __init__(self, uniform=None): + self.uniform = uniform + self.calls = [] + + def uniform_row_inputs(self, graph_walk): + return self.uniform + + def prepare_inputs(self, **kw): + self.calls.append(kw["graph_walk"]) + return object() + + def can_batch(self, batch, model_inputs): + return len(model_inputs) > 0 + + +def _engine(sub): + eng = Engine.__new__(Engine) + eng._submodules = {"llm": SimpleNamespace(submodule=sub, resources={})} + eng._enable_nvtx = False + return eng + + +def test_a_uniform_walk_is_prepared_with_one_object(): + row = object() + sub = _Sub(uniform=row) + batch = _Batch([1, 2, 3]) + _engine(sub).prepare_inputs(batch) + assert batch.inputs == [row, row, row] and all(x is row for x in batch.inputs) + assert sub.calls == [] and batch.running_batched is True + + +def test_other_walks_still_prepare_row_by_row(): + sub = _Sub(uniform=None) + batch = _Batch([1, 2]) + _engine(sub).prepare_inputs(batch) + assert sub.calls == ["decode", "decode"] and len(batch.inputs) == 2 From 9db4d7aafab06417c59bbcd6fa5276b1acd3ea59 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:25:27 -0700 Subject: [PATCH 173/219] qwen3_5: a captured decode step gathers its tokens and rotary tables in the graph --- mstar/model/qwen3_5/submodules.py | 51 +++++++++++++++++++++++++++---- 1 file changed, 45 insertions(+), 6 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index b1a1f0d77d..1f7d2ad3ae 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -313,6 +313,32 @@ def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: return frozenset({"text_inputs"}) return frozenset() + @staticmethod + def _captured_decode(engine_inputs: ModelInputsFromEngine) -> bool: + """A decode step replayed on a leased slot (or being captured for one).""" + step = getattr(engine_inputs, "step", None) + ctx = getattr(step, "ctx", None) if step is not None else None + return ( + ctx is not None + and getattr(ctx, "slot_lease", None) is not None + and getattr(ctx, "graph_walk", None) == "decode" + ) + + def _decode_cos_sin( + self, engine_inputs: ModelInputsFromEngine, num_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The rotary tables of a decode step from the position resource's + planned positions: one token per row, no custom grids.""" + pos_ids = self.node_resources[ROPE].pos_ids("main") + assert pos_ids is not None, ( + "position resource has no plan for this step; `plan` must run " + "before the forward" + ) + pos = pos_ids[:num_tokens].unsqueeze(0).expand(3, -1) + return self.model.model.build_cos_sin( + pos, dtype=self.model.model.embed_tokens.weight.dtype, + ) + def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: """``[3, total_tokens]`` for the step, in packed request order. @@ -352,6 +378,12 @@ def preprocess( ) -> dict[str, torch.Tensor | Any]: out: dict[str, torch.Tensor | Any] = {} if inputs[0].input_ids is None and inputs[0].input_embeds is None: + if self._captured_decode(engine_inputs): + # A captured decode step: the token gather and the rotary + # tables run inside the graph (``_forward``), off the static + # buffers the plan and the previous step's commit wrote, so + # the replay stages no input at all. + return out # device loop-back: every row of a decode step is its request's # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] @@ -432,13 +464,20 @@ def _forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor, - sin_3d: torch.Tensor, + cos_3d: torch.Tensor | None = None, + sin_3d: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, ) -> torch.Tensor: attn: AttentionManager = engine_inputs.resources[ATTN] sampler: SamplerResource = engine_inputs.resources[SAMPLER] + if input_ids is None and input_embeds is None: + # a captured decode step (see ``preprocess``): the gather off the + # sampler's slot master is recorded in the graph, so a replay reads + # the tokens the previous step's commit scattered there + input_ids = sampler.loopback_tokens(engine_inputs.request_ids) + if cos_3d is None or sin_3d is None: + cos_3d, sin_3d = self._decode_cos_sin(engine_inputs, input_ids.shape[0]) embeds = ( self.model.model.embed_tokens(input_ids) @@ -457,8 +496,8 @@ def forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor, - sin_3d: torch.Tensor, + cos_3d: torch.Tensor | None = None, + sin_3d: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, @@ -484,8 +523,8 @@ def forward_batched( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor, - sin_3d: torch.Tensor, + cos_3d: torch.Tensor | None = None, + sin_3d: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, From d16a5b6b82af77abb1f2270f4b48b54ee33eecf5 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:25:27 -0700 Subject: [PATCH 174/219] test: captured decode stages no input --- test/modular/test_qwen3_5_device_loopback.py | 46 ++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py index d98b282eb9..3bedc46500 100644 --- a/test/modular/test_qwen3_5_device_loopback.py +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -91,3 +91,49 @@ def test_the_loop_back_decode_row_is_uniform(monkeypatch): assert LLMSubmodule.uniform_row_inputs(_sub(_Sampler(masters=False)), "decode") is None monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") assert LLMSubmodule.uniform_row_inputs(sub, "decode") is None + + +def test_a_captured_decode_step_stages_no_input(): + """Under a leased decode step the gather and the rotary tables are part + of the graph, so preprocess hands the replay nothing to stage; an eager + step (no lease) still gets them from preprocess.""" + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4])) + sub = _sub(sampler) + rows = [_prepare(sub, "decode", {}) for _ in range(2)] + leased = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2"], + step=SimpleNamespace(ctx=SimpleNamespace(slot_lease=object(), graph_walk="decode")), + ) + assert LLMSubmodule.preprocess(sub, "decode", leased, rows) == {} + + eager = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2"], + step=SimpleNamespace(ctx=SimpleNamespace(slot_lease=None, graph_walk="decode")), + ) + sub.model = SimpleNamespace(model=SimpleNamespace( + build_cos_sin=lambda pos, dtype: (pos, pos), + embed_tokens=SimpleNamespace(weight=torch.zeros(1)), + )) + sub._position_ids_3d = lambda inputs: torch.zeros(3, len(inputs)) + out = LLMSubmodule.preprocess(sub, "decode", eager, rows) + assert out["input_ids"].tolist() == [3, 4] and "cos_3d" in out + + +def test_decode_cos_sin_comes_from_the_planned_positions(): + from mstar.model.qwen3_5.config import ROPE + + sub = LLMSubmodule.__new__(LLMSubmodule) + seen = {} + + def build(pos, dtype): + seen["pos"] = pos + return pos, pos + + sub.node_resources = {ROPE: SimpleNamespace(pos_ids=lambda label: torch.arange(8))} + sub.model = SimpleNamespace(model=SimpleNamespace( + build_cos_sin=build, + embed_tokens=SimpleNamespace(weight=torch.zeros(1, dtype=torch.bfloat16)), + )) + cos, sin = LLMSubmodule._decode_cos_sin(sub, SimpleNamespace(), 5) + assert seen["pos"].shape == (3, 5) and seen["pos"][0].tolist() == [0, 1, 2, 3, 4] + assert cos is seen["pos"] and sin is seen["pos"] From 5adfacc18204b3480250802fdab7626a64966765 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:29:02 -0700 Subject: [PATCH 175/219] knobs: MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT --- mstar/utils/knobs.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index ec5ef92f3e..38bc8dd695 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -8,3 +8,14 @@ def device_loopback_enabled() -> bool: routing a tensor per request; also whether the sampler persists the last token at all.""" return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" + + +def launch_signal_after_commit() -> bool: + """``MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT`` (default 0): release the thread + that submitted a step only once the gpu thread has also committed the + step and staged its outputs, instead of at the replay launch. Every torch + call after the launch releases the GIL, and the woken main thread keeps + it for its whole postprocess stretch, so the first such call waits that + long; signalling later costs the main thread the uncontended commit and + collect (a fraction of a millisecond) instead.""" + return os.environ.get("MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT", "0") == "1" From 3932a4969c186c985ea97fcb215415f151cb46d3 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:29:03 -0700 Subject: [PATCH 176/219] worker: _exec_step, launch signal after commit behind the knob --- mstar/worker/worker.py | 49 +++++++++++++++++++++++++++++++----------- 1 file changed, 37 insertions(+), 12 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 1d57d67fb2..59516eff73 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -73,6 +73,7 @@ WorkerMessage, WorkerMessageType, ) +from mstar.utils.knobs import launch_signal_after_commit from mstar.utils.profiler import PHASE_PERIOD, nvtx_enabled, phase_buffer, range_pop, range_push from mstar.worker.engine_manager import EngineManager from mstar.worker.micro_scheduler import MicroScheduler, ScheduledBatch @@ -295,6 +296,8 @@ class Worker: # Empty for a text model, which lets a speculation build skip its three # per-row stream polls; None (a bare instance) means unknown, so poll. _stream_buffer_rids: set[int] | None = None + # a bare instance keeps the launch signal where the engine sets it + _launch_signal_after_commit: bool = False def __init__( self, @@ -324,6 +327,7 @@ def __init__( self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} self._host_rows_slot = 0 self._stream_buffer_rids = set() + self._launch_signal_after_commit = launch_signal_after_commit() # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -1917,18 +1921,12 @@ def _execute_on_gpu_thread( self._phase_record( f"rows.{batch.graph_walk}#", len(node_batch.request_ids), ) - with self._span("worker.gpu_thread.exec"): - outputs = engine.exec_and_postprocess(node_batch) - if execution_stream is not None: - self._stage_host_rows(outputs) - event = torch.Event() - event.record(execution_stream) - node_batch.completion_event = event - if gpu_start is not None: - gpu_end = torch.Event(enable_timing=True) - gpu_end.record(execution_stream) - self._gpu_spans.append((gpu_start, gpu_end)) - self._drain_gpu_spans() + outputs = self._exec_step(engine, node_batch, execution_stream) + if gpu_start is not None: + gpu_end = torch.Event(enable_timing=True) + gpu_end.record(execution_stream) + self._gpu_spans.append((gpu_start, gpu_end)) + self._drain_gpu_spans() return outputs finally: # Safety net: a step that raised before the forward would otherwise @@ -3316,6 +3314,33 @@ def _inline_emit_values( return None, [] return None, [] + def _exec_step(self, engine, node_batch: ExecutingBatch, execution_stream): + """The step on the gpu thread: the engine's forward and commit, the + staged host rows, the completion event. + + With ``MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT`` the engine sees no event to + set at the replay launch; the submitter is released here afterwards, + once commit, collect and the staging are done, so none of those torch + calls has to win the GIL back from the thread it woke. + """ + deferred_signal = None + if self._launch_signal_after_commit and node_batch.launch_started_event is not None: + deferred_signal = node_batch.launch_started_event + node_batch.launch_started_event = None + try: + with self._span("worker.gpu_thread.exec"): + outputs = engine.exec_and_postprocess(node_batch) + if execution_stream is not None: + self._stage_host_rows(outputs) + event = torch.Event() + event.record(execution_stream) + node_batch.completion_event = event + finally: + if deferred_signal is not None: + node_batch.launch_started_event = deferred_signal + deferred_signal.set() + return outputs + def _stage_host_rows(self, outputs: "BatchedModelOutput") -> None: """Start the device-to-host copy of the step's row-addressed stop buffers here on the gpu thread, on the execution stream right behind From 999eab1c4c556c27748429c608fb6a3c1af42e1f Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:29:03 -0700 Subject: [PATCH 177/219] docs: MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT --- docs/environment_variables.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 09bb9500a3..bcc1e8b385 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -360,6 +360,15 @@ Worker scheduling 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_LAUNCH_SIGNAL_AFTER_COMMIT`` + - ``0`` + - Release the thread that submitted a step only once the gpu thread + has also committed it and staged its outputs, instead of at the + replay launch. Every torch call after the launch releases the GIL and + the woken main thread keeps it for its whole postprocess stretch, so + the first such call waits that long; signalling later costs the main + thread the uncontended commit and collect instead. 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 From f6e0d22ecbcc696b379804a71759df58ba43f6a3 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:29:03 -0700 Subject: [PATCH 178/219] test: launch signal placement --- test/modular/test_worker_launch_signal.py | 67 +++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 test/modular/test_worker_launch_signal.py diff --git a/test/modular/test_worker_launch_signal.py b/test/modular/test_worker_launch_signal.py new file mode 100644 index 0000000000..1c470e59f9 --- /dev/null +++ b/test/modular/test_worker_launch_signal.py @@ -0,0 +1,67 @@ +"""The submitter's launch signal: set by the engine at the replay launch by +default, or by the gpu thread after commit and staging when the knob is on.""" + +import threading +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from mstar.model.submodule_base import BatchedModelOutput # noqa: E402 +from mstar.worker.worker import Worker # noqa: E402 + + +class _Engine: + def __init__(self): + self.seen_event = "unset" + + def exec_and_postprocess(self, node_batch): + # what the engine would hand to the replay as its launch signal + self.seen_event = node_batch.launch_started_event + return BatchedModelOutput() + + +class _Span: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _worker(deferred): + w = object.__new__(Worker) + w._launch_signal_after_commit = deferred + w._span = lambda name: _Span() + return w + + +def test_by_default_the_engine_gets_the_event_to_set_at_the_launch(): + engine = _Engine() + event = threading.Event() + batch = SimpleNamespace(launch_started_event=event, completion_event=None) + Worker._exec_step(_worker(False), engine, batch, None) + assert engine.seen_event is event + assert not event.is_set(), "the stub engine never launched, so nobody set it" + + +def test_with_the_knob_the_gpu_thread_signals_after_the_step(): + engine = _Engine() + event = threading.Event() + batch = SimpleNamespace(launch_started_event=event, completion_event=None) + Worker._exec_step(_worker(True), engine, batch, None) + assert engine.seen_event is None, "the engine must not signal at the launch" + assert event.is_set() and batch.launch_started_event is event + + +def test_a_raising_step_still_releases_the_submitter(): + class _Raising(_Engine): + def exec_and_postprocess(self, node_batch): + raise RuntimeError("forward failed") + + event = threading.Event() + batch = SimpleNamespace(launch_started_event=event, completion_event=None) + with pytest.raises(RuntimeError): + Worker._exec_step(_worker(True), _Raising(), batch, None) + assert event.is_set() and batch.launch_started_event is event From b9412aaca89c3c031b4fa71b46ec0b5e3a3dbc94 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:32:35 -0700 Subject: [PATCH 179/219] qwen3_5: cheaper batched stop check --- mstar/model/qwen3_5/submodules.py | 32 +++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 1f7d2ad3ae..ecdd134796 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -599,23 +599,31 @@ def check_stop_batched( if not torch.is_tensor(tokens) or tokens.dim() == 0: return None values = tokens.reshape(tokens.shape[0], -1)[:, 0].tolist() - row_of = {rid: i for i, rid in enumerate(host_rows.request_ids)} + rows = host_rows.request_ids + if len(rows) >= len(request_ids) and all( + a == b for a, b in zip(request_ids, rows, strict=False) + ): + # the common case: the batch in forward order, no row map needed + pairs = zip(request_ids, values, strict=False) + else: + row_of = {rid: i for i, rid in enumerate(rows)} + pairs = ( + (rid, values[i]) for rid in request_ids + if (i := row_of.get(rid)) is not None and i < len(values) + ) # no row, no output: the per-request path skips it too stop_ids = self.config.stop_token_ids stops: dict[str, set[str]] = {} - for rid in request_ids: - i = row_of.get(rid) - if i is None or i >= len(values): - continue # no row, no output: the per-request path skips it too + for rid, value in pairs: info = request_infos[rid] - hit_eos = ( - not info.resource_configs[SAMPLER].ignore_eos - and values[i] in stop_ids - ) - out_of_budget = ( + # the set test first: a stop id is rare, the config read behind it + # goes through the wrapper + if ( + value in stop_ids + and not info.resource_configs[SAMPLER].ignore_eos + ) or ( info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 >= info.max_tokens - ) - if hit_eos or out_of_budget: + ): stops[rid] = {"decode_loop"} return stops From 6f83e1042e8551b7e06052fa26cdf4be5571d902 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:32:35 -0700 Subject: [PATCH 180/219] kv: the chain extension looks a request's label up once --- mstar/engine/resources/kv/manager.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 1204b7513d..3bd6026b42 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -365,6 +365,9 @@ def __init__( ) self._streams: dict[str, LabelToStream] = {} self._overrides: dict[str, KVReqConfig] = {} + # (node, walk) -> (label, stop buffer name) per request, for the + # decode-step chain extension; dropped with the request + self._chain_lookup: dict[str, dict[tuple[str, str], tuple]] = {} # its entity id is the worker id, and one worker is one copy of a node self._replica = ( transfer_engine_info.my_entity_id @@ -718,13 +721,26 @@ def extend_prefix_chains_batch( buffers = host_rows.buffers rows_of: dict[str, list] = {} page_size = self.config.page_size + key = (node_name, graph_walk) + lookup = self._chain_lookup with self._lock: for rid in request_ids: - label = self._keyed_label(rid, node_name, graph_walk) - if label is None: - continue - tensor = (self._overrides[rid].prefix_decode or {}).get(label) - if not tensor: + # the label and buffer name do not change over a request's + # life (its overrides are set at ingest), so each is looked up + # once instead of on every decode step + per_rid = lookup.get(rid) + if per_rid is None: + per_rid = lookup[rid] = {} + entry = per_rid.get(key) + if entry is None: + label = self._keyed_label(rid, node_name, graph_walk) + tensor = ( + (self._overrides[rid].prefix_decode or {}).get(label) + if label is not None else None + ) + entry = per_rid[key] = (label, tensor) + label, tensor = entry + if label is None or not tensor: continue rows = rows_of.get(tensor) if rows is None: @@ -2000,6 +2016,7 @@ def remove_request(self, rid: str): self._forget_reservations(rid) self._streams.pop(rid, None) self._overrides.pop(rid, None) + self._chain_lookup.pop(rid, None) self._transfer.remove_request(rid) if _DEBUG_ASSERTS: self.assert_pages_conserved() From e14a743b4afaab272b60395354902b417c8966fc Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:32:35 -0700 Subject: [PATCH 181/219] test: stop check off a reordered batch --- .../test_qwen3_5_check_stop_batched.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/test/modular/test_qwen3_5_check_stop_batched.py b/test/modular/test_qwen3_5_check_stop_batched.py index 7822615118..87e89994fe 100644 --- a/test/modular/test_qwen3_5_check_stop_batched.py +++ b/test/modular/test_qwen3_5_check_stop_batched.py @@ -88,3 +88,30 @@ def test_engine_falls_back_and_attributes_failures(): ) assert stops == {"a": {"decode_loop"}} assert list(failures) == ["x"] + + +def test_rows_out_of_batch_order_still_map_by_request(): + """A batch listed in another order than the forward ran it reads each + request's own row, and a request without a row is skipped.""" + from types import SimpleNamespace + + from mstar.model.qwen3_5.config import SAMPLER + from mstar.model.qwen3_5.submodules import LLMSubmodule + from mstar.model.submodule_base import HostRows + + sub = LLMSubmodule.__new__(LLMSubmodule) + sub.config = SimpleNamespace(stop_token_ids={7}) + rows = HostRows(("a", "b", "c"), {"new_token": torch.tensor([1, 7, 1])}) + + def info(max_tokens, iters=0): + return SimpleNamespace( + resource_configs={SAMPLER: SimpleNamespace(ignore_eos=False)}, + dynamic_loop_iter_counts={"decode_loop": iters}, max_tokens=max_tokens, + ) + + infos = {"a": info(100), "b": info(100), "c": info(5, iters=4), "d": info(100)} + stops = LLMSubmodule.check_stop_batched(sub, ["c", "d", "b", "a"], infos, rows) + assert stops == {"c": {"decode_loop"}, "b": {"decode_loop"}} + # forward order: the same answer without the row map + stops = LLMSubmodule.check_stop_batched(sub, ["a", "b", "c"], infos, rows) + assert stops == {"b": {"decode_loop"}, "c": {"decode_loop"}} From 6799af392e30d4c696b8f4ef95555b1fc383b6f6 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:36:38 -0700 Subject: [PATCH 182/219] api: inline tokens as ids for models that decode from ids --- mstar/api_server/data_worker.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index be0e70c379..ced6a5b434 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -737,9 +737,15 @@ def _emit_inline_token(self, item: InlineToken) -> None: sequence = state.next_sequence state.next_sequence += 1 try: - tensor = torch.tensor([item.value], dtype=torch.int64) + # A model that detokenizes from ids takes the plain list: one + # tensor build and one tensor-to-list per token saved, at tens of + # thousands of tokens a second through this thread. + if getattr(self.model, "inline_postprocess_takes_ids", False): + output = [item.value] + else: + output = torch.tensor([item.value], dtype=torch.int64) postprocessed = self.model.postprocess( - tensor, item.modality, + output, item.modality, request_kwargs=self.request_model_kwargs.get(request_id), ) self._queue_completed_output( From 44e7ba7205be3100d191499c6c2b5b019cb61f43 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:36:39 -0700 Subject: [PATCH 183/219] qwen3_5: inline postprocess takes ids --- mstar/model/qwen3_5/qwen3_5_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index ddd36b9160..e2036afa6e 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -443,9 +443,12 @@ def process_prompt( return result + # inline tokens arrive as ids; the tokenizer decodes a list as it does a tensor + inline_postprocess_takes_ids = True + def postprocess( self, - output: torch.Tensor, + output: "torch.Tensor | list[int]", modality: str, request_kwargs: dict | None = None, ) -> bytes: From 2d91a8d604a2e53e024cc2c91994702745f9fb73 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:36:39 -0700 Subject: [PATCH 184/219] test: inline value as a list --- test/modular/test_inline_tokens.py | 39 ++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py index a20cb35938..0e8292af8d 100644 --- a/test/modular/test_inline_tokens.py +++ b/test/modular/test_inline_tokens.py @@ -190,3 +190,42 @@ class _NoInline: assert _worker(_NoInline())._inline_emit_values( eng, _Batch(), rows, [7, 8], ["text_inputs"], ) == none + + +def test_a_model_that_decodes_ids_gets_the_inline_value_as_a_list(): + """The per-token tensor build is skipped for a model that opts in; the + default model still receives a tensor.""" + from mstar.api_server.data_worker import InlineToken, PreprocessWorkerThread + + class _Ids(_Model): + inline_postprocess_takes_ids = True + + def postprocess(self, output, modality, request_kwargs=None): + assert isinstance(output, list) and output == [5] + return b"<5>" + + def _worker(model): + w = object.__new__(PreprocessWorkerThread) + w.model = model + w.request_output_state = {} + w.request_model_kwargs = {} + w.queued = [] + w._queue_completed_output = lambda rid, seq, chunk: w.queued.append((rid, seq, chunk.data)) + w._fail_request = lambda *a, **kw: w.queued.append(("fail", a)) + return w + + item = InlineToken(request_id="r1", value=5, loop_indices=None, modality="text") + w = _worker(_Ids()) + w._emit_inline_token(item) + assert w.queued == [("r1", 0, b"<5>")] + + seen = {} + + class _Tensor(_Model): + def postprocess(self, output, modality, request_kwargs=None): + seen["type"] = type(output).__name__ + return b"x" + + w = _worker(_Tensor()) + w._emit_inline_token(item) + assert seen["type"] == "Tensor" and w.queued == [("r1", 0, b"x")] From 85bc776c4e0f5305f171edce383eb9a4f82b18d8 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:46:59 -0700 Subject: [PATCH 185/219] knobs: MSTAR_SAMPLER_INGRAPH_SCATTER, MSTAR_SAMPLER_SLOTS --- mstar/utils/knobs.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 38bc8dd695..7d6831d391 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -19,3 +19,21 @@ def launch_signal_after_commit() -> bool: long; signalling later costs the main thread the uncontended commit and collect (a fraction of a millisecond) instead.""" return os.environ.get("MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT", "0") == "1" + + +def sampler_ingraph_scatter() -> bool: + """``MSTAR_SAMPLER_INGRAPH_SCATTER`` (default 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 step + issues none of those launches on the gpu thread (each is a GIL release + the main thread's postprocess can win). Padding rows then address a trash + master row, and the masters are allocated once for + ``MSTAR_SAMPLER_SLOTS`` (default 4096) concurrent requests, since a + captured graph holds their addresses.""" + return os.environ.get("MSTAR_SAMPLER_INGRAPH_SCATTER", "0") == "1" + + +def sampler_slots() -> int: + """``MSTAR_SAMPLER_SLOTS`` (default 4096): master rows allocated up front + when the in-graph scatter is on (see ``sampler_ingraph_scatter``).""" + return int(os.environ.get("MSTAR_SAMPLER_SLOTS", "4096")) From 38c240ee4a0e1d935cda6791ca80fefd75d1615e Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:46:59 -0700 Subject: [PATCH 186/219] sampler: offset gather and the two commit scatters inside the graph, behind the knob --- mstar/engine/resources/sampler/utils.py | 68 ++++++++++++++++++++++--- 1 file changed, 62 insertions(+), 6 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index 7e3b5bff19..ccfb2f8d67 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -28,6 +28,7 @@ import triton.language as tl from mstar.utils.h2d import PinnedStager +from mstar.utils.knobs import sampler_slots logger = logging.getLogger(__name__) @@ -1030,16 +1031,40 @@ class CudaGraphableSampler(BaseSampler): # the slot master so the next step can read the token on the device. last_token_buf: torch.Tensor | None = None tp_group: "CommGroup | None" = None # noqa: F821 + # In-graph gather/scatter (``MSTAR_SAMPLER_INGRAPH_SCATTER``): the slot + # masters and this slot's padded index row, all static, so the captured + # sample reads its offsets from the masters and writes the advanced + # offsets and the sampled tokens back without a launch after the replay. + offset_master: torch.Tensor | None = None + last_token_master: torch.Tensor | None = None + slot_idx_view: torch.Tensor | None = None + ingraph_scatter: bool = False # Set during graph capture, and used by the cuda graph runner to determine # whether requests' seen token buffers should be synced post-replay applied_penalty_in_graph: bool = False + def gather_in_graph(self) -> None: + """Offsets of this step's rows off the master, inside the graph.""" + torch.index_select( + self.offset_master, 0, self.slot_idx_view, out=self.offset_buf, + ) + + def scatter_in_graph(self, codes: torch.Tensor) -> None: + """Advanced offsets and sampled tokens back to the masters, inside the + graph: every row, padding rows into the trash row.""" + self.offset_master.index_copy_(0, self.slot_idx_view, self.offset_buf) + self.last_token_master.index_copy_( + 0, self.slot_idx_view, codes.to(self.last_token_master.dtype), + ) + @torch.compiler.disable def sample( self, request_ids: list[str], logits: torch.Tensor, apply_penalty: bool = False, ): + if self.ingraph_scatter: + self.gather_in_graph() codes = sample_cuda_graphable_gpu( logits, self.temperature_buf, self.top_k_buf, self.top_p_buf, @@ -1053,6 +1078,8 @@ def sample( codes = self._broadcast_tokens(codes) if self.last_token_buf is not None: self.last_token_buf.copy_(codes) + if self.ingraph_scatter: + self.scatter_in_graph(codes) if apply_penalty and self.seen_tokens_buf is not None: self.applied_penalty_in_graph = True # Record the (broadcast, TP-agreed) token in the seen-token buffer so @@ -1311,6 +1338,10 @@ class SamplerBuffers: # Master cache capacity (grown by doubling when more requests are # concurrently registered than the per-step buffer holds). _master_capacity: int = field(default=0, repr=False) + # the in-graph scatter (knob): padding rows address ``_trash_slot``, one + # master row past the capacity that no request is ever given + ingraph_scatter: bool = False + _trash_slot: int = 0 # Number of double-buffer slots (per-step buffers carry this leading dim). cg_slots: int = 1 # Per-step slot-index staging, per cg slot. ``_slot_idx_cpu`` is pinned so @@ -1376,9 +1407,15 @@ def allocate( vocab_size: int | None = None, cg_slots: int = 1, enable_min_p: bool = False, + ingraph_scatter: bool = False, + slots: int | None = None, ) -> "SamplerBuffers": """Allocate sampling buffers for ``max_batch_size``. + With ``ingraph_scatter`` the masters hold ``slots`` rows (the knob's + 4096 by default) plus one trash row for padding, allocated once: the + captured graphs keep their addresses, so they never grow. + ``vocab_size`` (when not None) enables the seen-token mask buffer for the repetition penalty. The master rows default to a ``SamplingConfig()`` row (temp=1, top_k=0, top_p=1, rep_penalty=1) — what an unregistered slot @@ -1387,16 +1424,21 @@ def allocate( """ pinned = torch.cuda.is_available() and device.type == "cuda" cap = max_batch_size + if ingraph_scatter: + cap = max(max_batch_size, slots if slots is not None else sampler_slots()) + # one row past the capacity is the trash row for padding when the + # in-graph scatter is on (harmless, never read for a request) + rows = cap + 1 if ingraph_scatter else cap def mk(dtype: torch.dtype, default: float) -> HostBuffer: return HostBuffer.allocate( - max_batch_size, cap, device, dtype, default, cg_slots + max_batch_size, rows, device, dtype, default, cg_slots ) # seen token mask is not double-buffered, as it depends on the GPU # value of the previous step. Same goes for offset seen_tokens = ( - MaskBuffer.allocate(max_batch_size, cap, vocab_size, device, cg_slots=1) + MaskBuffer.allocate(max_batch_size, rows, vocab_size, device, cg_slots=1) if vocab_size is not None else None ) return cls( @@ -1406,12 +1448,14 @@ def mk(dtype: torch.dtype, default: float) -> HostBuffer: top_p=mk(torch.float32, 1.0), seed=mk(torch.long, 0), rep_penalty=mk(torch.float32, 1.0), - offset=Buffer.allocate(max_batch_size, cap, device, torch.long, 0, 1), - last_token=Buffer.allocate(max_batch_size, cap, device, torch.long, 0, cg_slots), + offset=Buffer.allocate(max_batch_size, rows, device, torch.long, 0, 1), + last_token=Buffer.allocate(max_batch_size, rows, device, torch.long, 0, cg_slots), tp_group=tp_group, seen_tokens=seen_tokens, min_p=mk(torch.float32, 0.0) if enable_min_p else None, _master_capacity=cap, + ingraph_scatter=ingraph_scatter, + _trash_slot=cap if ingraph_scatter else 0, cg_slots=cg_slots, _slot_idx_cpu=torch.zeros(cg_slots, max_batch_size, dtype=torch.long, pin_memory=pinned), _slot_idx_gpu=torch.zeros(cg_slots, max_batch_size, dtype=torch.long, device=device), @@ -1435,6 +1479,10 @@ def slice_for_bs(self, bs: int, cg_slot: int = 0) -> dict[str, Any]: "min_p_buf": self.min_p.slot_view(cg_slot, bs) if self.min_p is not None else None, "last_token_buf": self.last_token.slot_view(cg_slot, bs), "tp_group": self.tp_group, + "offset_master": self.offset.master, + "last_token_master": self.last_token.master, + "slot_idx_view": self._slot_idx_gpu[cg_slot, :bs], + "ingraph_scatter": self.ingraph_scatter, } # ------------------------------------------------------------------ @@ -1474,6 +1522,12 @@ def _grow_master(self, new_capacity: int) -> None: master capacity. Per-step buffers (sized to the cuda-graph max_bs) are NOT resized — the gather only reads ``padded_bs`` rows from master. """ + if self.ingraph_scatter: + raise RuntimeError( + f"more than {self._master_capacity} concurrent requests on the " + "sampler with the in-graph scatter on; raise MSTAR_SAMPLER_SLOTS " + "or set MSTAR_SAMPLER_INGRAPH_SCATTER=0" + ) for buf in self._scalar_buffers(): buf.grow_master(new_capacity) self.offset.grow_master(new_capacity) @@ -1610,7 +1664,7 @@ def _stage_slot_idx( self._init_slot(slot, rid) rows.append(slot) if len(rows) < padded_bs: - rows.extend([0] * (padded_bs - len(rows))) + rows.extend([self._trash_slot] * (padded_bs - len(rows))) self._slot_idx_np[cg_slot, :padded_bs] = rows self._staged_rids[cg_slot] = (tuple(request_ids), padded_bs) @@ -1673,7 +1727,9 @@ def gather_dynamic( for slot in resets: self.offset.master[slot:slot + 1].zero_() idx_view = self._upload_slot_idx(padded_bs, cg_slot) - self.offset.gather(idx_view, padded_bs, cg_slot) + if not self.ingraph_scatter: + # with it on, the captured sample gathers off the master itself + self.offset.gather(idx_view, padded_bs, cg_slot) if self.seen_tokens is not None and gather_seen_tokens: self.seen_tokens.gather(idx_view, padded_bs, cg_slot) From 8f6e60e27049641d6f0622e31a57467649a5282b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:47:00 -0700 Subject: [PATCH 187/219] sampler: commit skips the scatters the graph did --- mstar/engine/resources/sampler/resource.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index 0a4e688e0c..0bf7612b1a 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -14,7 +14,7 @@ ) from mstar.engine.resources.sampler.utils import CudaGraphableSampler, Sampler, SamplerBuffers from mstar.engine.resources.step import SlotLease, StepContext -from mstar.utils.knobs import device_loopback_enabled +from mstar.utils.knobs import device_loopback_enabled, sampler_ingraph_scatter from mstar.utils.profiler import PHASE_PERIOD, phase_record @@ -126,6 +126,7 @@ def build_cuda_graph_buffers( vocab_size=self._vocab_size, cg_slots=self._cg_slots, enable_min_p=self._enable_min_p, + ingraph_scatter=self._keep_last_token and sampler_ingraph_scatter(), ) def ingest_request(self, rid: str, overrides: SamplingReqConfig | None=None): @@ -321,6 +322,13 @@ def commit(self, step: SamplerStep, ctx: StepContext): # None on an eager step, which never gathered one if self._cg_buffers is None or self._cg_sampler is None: return + if self._cg_buffers.ingraph_scatter: + # the captured sample scattered both rows itself + if self._penalty_needed_this_step: + self._cg_sampler.sync_seen_token_masks( + [self._sampler.get_token_mask(rid) for rid in ctx.request_ids] + ) + return t0 = perf_counter() if PHASE_PERIOD else 0.0 self._cg_buffers.scatter_offset(ctx.slot_lease.slot) if PHASE_PERIOD: From 8bedee480af76f7ddd97ba03e7422dad99fe3f24 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:47:00 -0700 Subject: [PATCH 188/219] docs: MSTAR_SAMPLER_INGRAPH_SCATTER --- docs/environment_variables.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index bcc1e8b385..6a75a92b23 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -360,6 +360,16 @@ Worker scheduling 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_LAUNCH_SIGNAL_AFTER_COMMIT`` - ``0`` - Release the thread that submitted a step only once the gpu thread From 9cc1fdd9bb86689ab28f8aa93d9955229ad02fb2 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:47:00 -0700 Subject: [PATCH 189/219] test: in-graph scatter rows and the trash row --- test/modular/test_sampler_last_token.py | 49 +++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py index d2cb346744..408a7e5a35 100644 --- a/test/modular/test_sampler_last_token.py +++ b/test/modular/test_sampler_last_token.py @@ -2,6 +2,7 @@ captured sample writes a per-step row, commit scatters it to the slot master, the next step gathers its input ids by slot; an eager step writes and reads the master by request.""" +import pytest import torch from mstar.engine.resources.sampler.utils import SamplerBuffers @@ -59,3 +60,51 @@ def test_growth_and_slot_reuse_keep_the_tokens(): bufs.register_request("d") # reuses a's slot; its first step writes it bufs.write_last_tokens(["d"], torch.tensor([8])) assert bufs.last_tokens_for(["d", "b", "c"]).tolist() == [8, 6, 7] + + +def _ingraph(max_bs=4, cg_slots=2, slots=8): + return SamplerBuffers.allocate( + max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, + ingraph_scatter=True, slots=slots, + ) + + +def test_ingraph_scatter_rows_land_by_slot_and_padding_hits_the_trash_row(): + """What the captured sample does with the knob on, step by step on the + CPU views: gather the offsets off the master, advance, scatter offsets + and tokens back; padding rows read and write the trash row only.""" + bufs = _ingraph() + assert bufs.ingraph_scatter and bufs._trash_slot == 8 + assert bufs.last_token.master.shape[0] == 9 and bufs.offset.master.shape[0] == 9 + for rid in ("a", "b", "c"): + bufs.register_request(rid) + bufs.write_last_tokens(["a", "b", "c"], torch.tensor([1, 2, 3])) + bufs.offset.master[bufs._rid_to_slot["a"]] = 10 + bufs.gather_static(["b", "a"], 4, 1) + bufs.gather_dynamic(["b", "a"], 4, 1) # uploads the index row, no offset gather + assert bufs._slot_idx_gpu[1, :4].tolist()[2:] == [8, 8] + s = bufs.sampler_for(4, 1) + assert s.ingraph_scatter and s.slot_idx_view.tolist()[:2] == [bufs._rid_to_slot[r] for r in ("b", "a")] + s.gather_in_graph() + assert s.offset_buf.tolist()[1] == 10 + s.offset_buf += 1 + s.scatter_in_graph(torch.tensor([22, 21, 99, 98])) + assert bufs.last_tokens_for(["a", "b", "c"]).tolist() == [21, 22, 3] + assert bufs.offset.master[bufs._rid_to_slot["a"]].item() == 11 + assert bufs.last_token.master[8].item() in (99, 98) + # the next step's input ids: real rows by slot, padding rows the trash row + got = bufs.gather_last_tokens(1, 4).tolist() + assert got[:2] == [22, 21] and got[2] == got[3] == bufs.last_token.master[8].item() + + +def test_ingraph_scatter_masters_never_grow(): + bufs = _ingraph(max_bs=2, slots=2) + bufs.register_request("a") + bufs.register_request("b") + with pytest.raises(RuntimeError, match="MSTAR_SAMPLER_SLOTS"): + bufs.register_request("c") + # the default path still grows + plain = SamplerBuffers.allocate(max_batch_size=2, device=torch.device("cpu")) + for rid in ("a", "b", "c"): + plain.register_request(rid) + assert plain.last_token.master.shape[0] >= 3 and not plain.ingraph_scatter From ee4bdbd9621a21ca46247938223c1c135842c128 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 19:48:00 -0700 Subject: [PATCH 190/219] test: offset reset runs before the in-graph gather --- test/modular/test_sampler_last_token.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py index 408a7e5a35..1df7393116 100644 --- a/test/modular/test_sampler_last_token.py +++ b/test/modular/test_sampler_last_token.py @@ -79,9 +79,10 @@ def test_ingraph_scatter_rows_land_by_slot_and_padding_hits_the_trash_row(): for rid in ("a", "b", "c"): bufs.register_request(rid) bufs.write_last_tokens(["a", "b", "c"], torch.tensor([1, 2, 3])) - bufs.offset.master[bufs._rid_to_slot["a"]] = 10 bufs.gather_static(["b", "a"], 4, 1) bufs.gather_dynamic(["b", "a"], 4, 1) # uploads the index row, no offset gather + # (a fresh slot's offset reset ran there; an advanced offset survives) + bufs.offset.master[bufs._rid_to_slot["a"]] = 10 assert bufs._slot_idx_gpu[1, :4].tolist()[2:] == [8, 8] s = bufs.sampler_for(4, 1) assert s.ingraph_scatter and s.slot_idx_view.tolist()[:2] == [bufs._rid_to_slot[r] for r in ("b", "a")] From 9500c8f9aa85a44eef5b3f1c0c4a9f0d3f0d385b Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:17:47 -0700 Subject: [PATCH 191/219] Revert "test: inline value as a list" This reverts commit 2d91a8d604a2e53e024cc2c91994702745f9fb73. --- test/modular/test_inline_tokens.py | 39 ------------------------------ 1 file changed, 39 deletions(-) diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py index 0e8292af8d..a20cb35938 100644 --- a/test/modular/test_inline_tokens.py +++ b/test/modular/test_inline_tokens.py @@ -190,42 +190,3 @@ class _NoInline: assert _worker(_NoInline())._inline_emit_values( eng, _Batch(), rows, [7, 8], ["text_inputs"], ) == none - - -def test_a_model_that_decodes_ids_gets_the_inline_value_as_a_list(): - """The per-token tensor build is skipped for a model that opts in; the - default model still receives a tensor.""" - from mstar.api_server.data_worker import InlineToken, PreprocessWorkerThread - - class _Ids(_Model): - inline_postprocess_takes_ids = True - - def postprocess(self, output, modality, request_kwargs=None): - assert isinstance(output, list) and output == [5] - return b"<5>" - - def _worker(model): - w = object.__new__(PreprocessWorkerThread) - w.model = model - w.request_output_state = {} - w.request_model_kwargs = {} - w.queued = [] - w._queue_completed_output = lambda rid, seq, chunk: w.queued.append((rid, seq, chunk.data)) - w._fail_request = lambda *a, **kw: w.queued.append(("fail", a)) - return w - - item = InlineToken(request_id="r1", value=5, loop_indices=None, modality="text") - w = _worker(_Ids()) - w._emit_inline_token(item) - assert w.queued == [("r1", 0, b"<5>")] - - seen = {} - - class _Tensor(_Model): - def postprocess(self, output, modality, request_kwargs=None): - seen["type"] = type(output).__name__ - return b"x" - - w = _worker(_Tensor()) - w._emit_inline_token(item) - assert seen["type"] == "Tensor" and w.queued == [("r1", 0, b"x")] From 01a929efce16f6fe37bfef9e09d1433988f359cf Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:17:48 -0700 Subject: [PATCH 192/219] Revert "qwen3_5: inline postprocess takes ids" This reverts commit 44e7ba7205be3100d191499c6c2b5b019cb61f43. --- mstar/model/qwen3_5/qwen3_5_model.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/mstar/model/qwen3_5/qwen3_5_model.py b/mstar/model/qwen3_5/qwen3_5_model.py index e2036afa6e..ddd36b9160 100644 --- a/mstar/model/qwen3_5/qwen3_5_model.py +++ b/mstar/model/qwen3_5/qwen3_5_model.py @@ -443,12 +443,9 @@ def process_prompt( return result - # inline tokens arrive as ids; the tokenizer decodes a list as it does a tensor - inline_postprocess_takes_ids = True - def postprocess( self, - output: "torch.Tensor | list[int]", + output: torch.Tensor, modality: str, request_kwargs: dict | None = None, ) -> bytes: From 7d3b9fe634a1ec717b8ceb249d3d4346bcadb7cb Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:17:48 -0700 Subject: [PATCH 193/219] Revert "api: inline tokens as ids for models that decode from ids" This reverts commit 6799af392e30d4c696b8f4ef95555b1fc383b6f6. --- mstar/api_server/data_worker.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index ced6a5b434..be0e70c379 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -737,15 +737,9 @@ def _emit_inline_token(self, item: InlineToken) -> None: sequence = state.next_sequence state.next_sequence += 1 try: - # A model that detokenizes from ids takes the plain list: one - # tensor build and one tensor-to-list per token saved, at tens of - # thousands of tokens a second through this thread. - if getattr(self.model, "inline_postprocess_takes_ids", False): - output = [item.value] - else: - output = torch.tensor([item.value], dtype=torch.int64) + tensor = torch.tensor([item.value], dtype=torch.int64) postprocessed = self.model.postprocess( - output, item.modality, + tensor, item.modality, request_kwargs=self.request_model_kwargs.get(request_id), ) self._queue_completed_output( From fbc37e1d8b78f81391b483df223527e2744a7ef8 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:22:52 -0700 Subject: [PATCH 194/219] utils: the chunked encode turns the tokenizer crate's parallelism on --- mstar/utils/fast_tokenize.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/mstar/utils/fast_tokenize.py b/mstar/utils/fast_tokenize.py index 286f61af76..88e6fddca6 100644 --- a/mstar/utils/fast_tokenize.py +++ b/mstar/utils/fast_tokenize.py @@ -61,6 +61,18 @@ def split_at(text: str, cuts: list[int]) -> list[str]: return parts +def _parallelism_on() -> None: + """The Rust backend encodes a batch on its thread pool only while + ``TOKENIZERS_PARALLELISM`` reads true; serving environments often export + it false (to silence the fork warning), which makes the chunked encode + no faster than one encode (18 vs 19 ms for 8.8k tokens on H200 hosts). + ``MSTAR_TOKENIZE_PARALLEL=0`` leaves the variable alone.""" + if os.environ.get("MSTAR_TOKENIZE_PARALLEL", "1") != "1": + return + if os.environ.get("TOKENIZERS_PARALLELISM", "").lower() not in ("true", "1"): + os.environ["TOKENIZERS_PARALLELISM"] = "true" + + def encode_ids( tokenizer, text: str, chunk_chars: int = DEFAULT_CHUNK_CHARS, max_chunks: int = DEFAULT_MAX_CHUNKS, @@ -73,6 +85,7 @@ def encode_ids( cuts = cut_points(text, chunk_chars, max_chunks) if backend is not None else [] if not cuts: return list(tokenizer(text).input_ids) + _parallelism_on() encodings = backend.encode_batch(split_at(text, cuts), add_special_tokens=False) ids: list[int] = [] for enc in encodings: From 2b13bfa819829177d9766d86ad1b23bf587165ed Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:22:52 -0700 Subject: [PATCH 195/219] docs: MSTAR_TOKENIZE_PARALLEL --- docs/environment_variables.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 6a75a92b23..517ccabb78 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -237,6 +237,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 From d87e0346d3ac0c239d5a73c3d680f0b8b4e5febd Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:22:52 -0700 Subject: [PATCH 196/219] test: chunked encode flips TOKENIZERS_PARALLELISM --- test/modular/test_fast_tokenize.py | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/test/modular/test_fast_tokenize.py b/test/modular/test_fast_tokenize.py index df8b35bdab..1e4da0f51d 100644 --- a/test/modular/test_fast_tokenize.py +++ b/test/modular/test_fast_tokenize.py @@ -55,3 +55,35 @@ def test_chunked_ids_are_identical_to_one_encode_on_the_benchmark_prompts(): whole = tok(t).input_ids assert encode_ids(tok, t, chunk_chars=2048, max_chunks=8) == whole assert encode_ids(tok, t, chunk_chars=512, max_chunks=16) == whole + + +def test_chunked_encode_turns_the_crate_parallelism_on(monkeypatch): + """The Rust backend only runs a batch in parallel while + TOKENIZERS_PARALLELISM reads true; a serving environment that exported + it false gets it flipped before the batch, unless the knob says no.""" + from mstar.utils.fast_tokenize import encode_ids + + class _Enc: + def __init__(self, ids): + self.ids = ids + + class _Backend: + def encode_batch(self, parts, add_special_tokens=False): + return [_Enc([len(p)]) for p in parts] + + class _Tok: + backend_tokenizer = _Backend() + + def __call__(self, text): + raise AssertionError("the chunked path was expected") + + text = "word " * 4000 + monkeypatch.setenv("TOKENIZERS_PARALLELISM", "false") + monkeypatch.delenv("MSTAR_TOKENIZE_PARALLEL", raising=False) + ids = encode_ids(_Tok(), text, chunk_chars=4096, max_chunks=8) + assert sum(ids) == len(text) and os.environ["TOKENIZERS_PARALLELISM"] == "true" + + monkeypatch.setenv("TOKENIZERS_PARALLELISM", "false") + monkeypatch.setenv("MSTAR_TOKENIZE_PARALLEL", "0") + encode_ids(_Tok(), text, chunk_chars=4096, max_chunks=8) + assert os.environ["TOKENIZERS_PARALLELISM"] == "false" From b0e3cc837d1837a6f39b043bea0f0c208d5e2fa8 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:25:10 -0700 Subject: [PATCH 197/219] knobs: MSTAR_TP_EARLY_SPEC --- mstar/utils/knobs.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 7d6831d391..7d27dfe66d 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -37,3 +37,13 @@ def sampler_slots() -> int: """``MSTAR_SAMPLER_SLOTS`` (default 4096): master rows allocated up front when the in-graph scatter is on (see ``sampler_ingraph_scatter``).""" return int(os.environ.get("MSTAR_SAMPLER_SLOTS", "4096")) + + +def tp_early_spec() -> bool: + """``MSTAR_TP_EARLY_SPEC`` (default 0): let a tensor-parallel leader build + its next speculation early (during the current step's postprocess, + committing only once it holds a head). Off by default: with the device + loop-back on a 4-way group this stalled every request for ~1.2 s once + (27B TP4 c8 674 vs 1009 tok/s with it off), and the TP cells are GPU-bound + anyway. TP1 nodes build early regardless (``MSTAR_EARLY_SPEC``).""" + return os.environ.get("MSTAR_TP_EARLY_SPEC", "0") == "1" From 3363443f57a6981cdc23203a5f1a854b00548e76 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:25:11 -0700 Subject: [PATCH 198/219] worker: a parallel node builds early only with the knob --- mstar/worker/worker.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 59516eff73..3dcb8f9bb2 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -73,7 +73,7 @@ WorkerMessage, WorkerMessageType, ) -from mstar.utils.knobs import launch_signal_after_commit +from mstar.utils.knobs import launch_signal_after_commit, tp_early_spec from mstar.utils.profiler import PHASE_PERIOD, nvtx_enabled, phase_buffer, range_pop, range_push from mstar.worker.engine_manager import EngineManager from mstar.worker.micro_scheduler import MicroScheduler, ScheduledBatch @@ -3963,6 +3963,8 @@ def run(self) -> None: # marker and schedules no yield-away, and the top of the next # iteration decides as it always did. early_spec = os.environ.get("MSTAR_EARLY_SPEC", "1") == "1" + # a parallel node builds early only with the knob (see tp_early_spec) + early_spec_tp = tp_early_spec() next_speculation: Speculation | None = None next_yield_away: tuple[str, str] | None = None from mstar.utils.profiler import range_pop, range_push @@ -4281,6 +4283,11 @@ def _maybe_clear_spec(): if ( early_spec and pp_state is not None and spec_pending is not None + and ( + early_spec_tp + or spec_pending.node_name + not in self.parallel_nodes + ) ): # N's stops are decided and N+1 is launched: # build N+2 now and hand its pre-plan to the From 14287922928f1ee829ba992a8bcd830e02307fd7 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:25:11 -0700 Subject: [PATCH 199/219] docs: MSTAR_TP_EARLY_SPEC --- docs/environment_variables.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 517ccabb78..2ac2d9d29d 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -377,6 +377,14 @@ Worker scheduling 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_TP_EARLY_SPEC`` + - ``0`` + - Let a tensor-parallel leader build its next speculation early, during + the current step's postprocess, committing only once it holds a head + (``MSTAR_EARLY_SPEC`` covers single-GPU nodes). Off by default: with + the device loop-back on a 4-way group every request stalled once for + about 1.2 s (27B TP4 at c8: 674 vs 1009 tok/s with it off), and the + tensor-parallel cells are GPU-bound, so there is nothing to buy. * - ``MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT`` - ``0`` - Release the thread that submitted a step only once the gpu thread From bb758445dc8aceebdb92654ca053c3d4eb146d14 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 20:25:11 -0700 Subject: [PATCH 200/219] test: tp early build knob --- test/modular/test_tp_early_spec_decision.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/modular/test_tp_early_spec_decision.py b/test/modular/test_tp_early_spec_decision.py index f3ab1890f9..ec10f068d4 100644 --- a/test/modular/test_tp_early_spec_decision.py +++ b/test/modular/test_tp_early_spec_decision.py @@ -92,3 +92,12 @@ def test_a_tp1_early_build_decides_on_the_spot(): w._tp_lead_needs_marker = lambda pending, spec: False # not a parallel node spec, target = w._build_speculation(_pending(), 1, _fair(), decide=True) assert spec is None and w.markers == [] and w.scheduler.get_next_calls == 1 + + +def test_the_tp_early_build_is_behind_a_knob(monkeypatch): + from mstar.utils.knobs import tp_early_spec + + monkeypatch.delenv("MSTAR_TP_EARLY_SPEC", raising=False) + assert not tp_early_spec() + monkeypatch.setenv("MSTAR_TP_EARLY_SPEC", "1") + assert tp_early_spec() From a66c7ad16544ad5c241295e957e61796c591e275 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 22:50:51 -0700 Subject: [PATCH 201/219] Revert "test: captured decode stages no input" This reverts commit d16a5b6b82af77abb1f2270f4b48b54ee33eecf5. --- test/modular/test_qwen3_5_device_loopback.py | 46 -------------------- 1 file changed, 46 deletions(-) diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py index 3bedc46500..d98b282eb9 100644 --- a/test/modular/test_qwen3_5_device_loopback.py +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -91,49 +91,3 @@ def test_the_loop_back_decode_row_is_uniform(monkeypatch): assert LLMSubmodule.uniform_row_inputs(_sub(_Sampler(masters=False)), "decode") is None monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") assert LLMSubmodule.uniform_row_inputs(sub, "decode") is None - - -def test_a_captured_decode_step_stages_no_input(): - """Under a leased decode step the gather and the rotary tables are part - of the graph, so preprocess hands the replay nothing to stage; an eager - step (no lease) still gets them from preprocess.""" - sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4])) - sub = _sub(sampler) - rows = [_prepare(sub, "decode", {}) for _ in range(2)] - leased = SimpleNamespace( - resources={SAMPLER: sampler}, request_ids=["r1", "r2"], - step=SimpleNamespace(ctx=SimpleNamespace(slot_lease=object(), graph_walk="decode")), - ) - assert LLMSubmodule.preprocess(sub, "decode", leased, rows) == {} - - eager = SimpleNamespace( - resources={SAMPLER: sampler}, request_ids=["r1", "r2"], - step=SimpleNamespace(ctx=SimpleNamespace(slot_lease=None, graph_walk="decode")), - ) - sub.model = SimpleNamespace(model=SimpleNamespace( - build_cos_sin=lambda pos, dtype: (pos, pos), - embed_tokens=SimpleNamespace(weight=torch.zeros(1)), - )) - sub._position_ids_3d = lambda inputs: torch.zeros(3, len(inputs)) - out = LLMSubmodule.preprocess(sub, "decode", eager, rows) - assert out["input_ids"].tolist() == [3, 4] and "cos_3d" in out - - -def test_decode_cos_sin_comes_from_the_planned_positions(): - from mstar.model.qwen3_5.config import ROPE - - sub = LLMSubmodule.__new__(LLMSubmodule) - seen = {} - - def build(pos, dtype): - seen["pos"] = pos - return pos, pos - - sub.node_resources = {ROPE: SimpleNamespace(pos_ids=lambda label: torch.arange(8))} - sub.model = SimpleNamespace(model=SimpleNamespace( - build_cos_sin=build, - embed_tokens=SimpleNamespace(weight=torch.zeros(1, dtype=torch.bfloat16)), - )) - cos, sin = LLMSubmodule._decode_cos_sin(sub, SimpleNamespace(), 5) - assert seen["pos"].shape == (3, 5) and seen["pos"][0].tolist() == [0, 1, 2, 3, 4] - assert cos is seen["pos"] and sin is seen["pos"] From df647563f425bbc981369ecdf334095fcb8f1c59 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 8 Oct 2026 22:50:51 -0700 Subject: [PATCH 202/219] Revert "qwen3_5: a captured decode step gathers its tokens and rotary tables in the graph" This reverts commit 9db4d7aafab06417c59bbcd6fa5276b1acd3ea59. --- mstar/model/qwen3_5/submodules.py | 51 ++++--------------------------- 1 file changed, 6 insertions(+), 45 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index ecdd134796..32e5518aa6 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -313,32 +313,6 @@ def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: return frozenset({"text_inputs"}) return frozenset() - @staticmethod - def _captured_decode(engine_inputs: ModelInputsFromEngine) -> bool: - """A decode step replayed on a leased slot (or being captured for one).""" - step = getattr(engine_inputs, "step", None) - ctx = getattr(step, "ctx", None) if step is not None else None - return ( - ctx is not None - and getattr(ctx, "slot_lease", None) is not None - and getattr(ctx, "graph_walk", None) == "decode" - ) - - def _decode_cos_sin( - self, engine_inputs: ModelInputsFromEngine, num_tokens: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - """The rotary tables of a decode step from the position resource's - planned positions: one token per row, no custom grids.""" - pos_ids = self.node_resources[ROPE].pos_ids("main") - assert pos_ids is not None, ( - "position resource has no plan for this step; `plan` must run " - "before the forward" - ) - pos = pos_ids[:num_tokens].unsqueeze(0).expand(3, -1) - return self.model.model.build_cos_sin( - pos, dtype=self.model.model.embed_tokens.weight.dtype, - ) - def _position_ids_3d(self, inputs: list[ARNodeInputs]) -> torch.Tensor: """``[3, total_tokens]`` for the step, in packed request order. @@ -378,12 +352,6 @@ def preprocess( ) -> dict[str, torch.Tensor | Any]: out: dict[str, torch.Tensor | Any] = {} if inputs[0].input_ids is None and inputs[0].input_embeds is None: - if self._captured_decode(engine_inputs): - # A captured decode step: the token gather and the rotary - # tables run inside the graph (``_forward``), off the static - # buffers the plan and the previous step's commit wrote, so - # the replay stages no input at all. - return out # device loop-back: every row of a decode step is its request's # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] @@ -464,20 +432,13 @@ def _forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor | None = None, - sin_3d: torch.Tensor | None = None, + cos_3d: torch.Tensor, + sin_3d: torch.Tensor, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, ) -> torch.Tensor: attn: AttentionManager = engine_inputs.resources[ATTN] sampler: SamplerResource = engine_inputs.resources[SAMPLER] - if input_ids is None and input_embeds is None: - # a captured decode step (see ``preprocess``): the gather off the - # sampler's slot master is recorded in the graph, so a replay reads - # the tokens the previous step's commit scattered there - input_ids = sampler.loopback_tokens(engine_inputs.request_ids) - if cos_3d is None or sin_3d is None: - cos_3d, sin_3d = self._decode_cos_sin(engine_inputs, input_ids.shape[0]) embeds = ( self.model.model.embed_tokens(input_ids) @@ -496,8 +457,8 @@ def forward( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor | None = None, - sin_3d: torch.Tensor | None = None, + cos_3d: torch.Tensor, + sin_3d: torch.Tensor, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, @@ -523,8 +484,8 @@ def forward_batched( self, graph_walk: str, engine_inputs: ModelInputsFromEngine, - cos_3d: torch.Tensor | None = None, - sin_3d: torch.Tensor | None = None, + cos_3d: torch.Tensor, + sin_3d: torch.Tensor, input_ids: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, From e3c54f55b4680ec0689ebee7947ca5dff1b543ed Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 203/219] engine: a reset pre-plan drops its kept step --- mstar/engine/engine.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 4e54522116..cce7bd7d71 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -1901,10 +1901,12 @@ def reset_pre_plan_for_batch(self, batch: ExecutingBatch | None = None) -> None: mark(f"engine.preplan_discarded.{batch.node_name}") batch.preplan_event = None batch.preplanned_rids = None - if batch.inputs is None: - # declared over template rows for a step that is now stale: - # exec declares again over the real ones - batch.step = None + # The step was declared over the rows the pre-plan saw: template + # rows for an unprepared batch, or a request set that has since + # lost rows to prepare_inputs (preplan_is_stale). Either way exec + # declares and admits again over the batch as it stands now; a + # kept step would go to admit unchecked. + batch.step = None lease = batch.step_context.slot_lease cg_runner = self._submodules[batch.node_name].cuda_graph_runner if lease is not None and cg_runner is not None: From f6274b1fa795c91545b0c62f94a799c986c155bf Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 204/219] test: a partly dropped batch does not keep its pre-planned step --- test/modular/test_preplan_error_paths.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/modular/test_preplan_error_paths.py b/test/modular/test_preplan_error_paths.py index db9e9a1467..da26c00d07 100644 --- a/test/modular/test_preplan_error_paths.py +++ b/test/modular/test_preplan_error_paths.py @@ -191,3 +191,26 @@ def test_exec_leaves_a_promoted_or_absent_stage_alone() -> None: with pytest.raises(ValueError): _exec_fake(runner, ValueError("forward failed")).exec(_batch(preplanned=None)) assert cleared == [] + + +def test_reset_drops_the_kept_step_of_a_partly_dropped_batch() -> None: + # pre-planned over ["a", "b"] with inputs in hand (the step was kept for + # exec), then prepare_inputs dropped "b": the stale step must not reach + # admit, so the reset clears it even though the batch still has inputs + cleared = [] + fake = SimpleNamespace( + _enable_nvtx=False, + _submodules={"dit": SimpleNamespace(cuda_graph_runner=None)}, + _runner=SimpleNamespace(clear_preplan=lambda: cleared.append(True)), + ) + fake.reset_pre_plan_for_batch = Engine.reset_pre_plan_for_batch.__get__(fake) + fake.preplan_is_stale = Engine.preplan_is_stale.__get__(fake) + batch = SimpleNamespace( + request_ids=["a"], node_name="dit", inputs=[object()], step=object(), + preplanned_rids=("a", "b"), preplan_event=object(), + step_context=SimpleNamespace(slot_lease=None), + ) + assert fake.preplan_is_stale(batch) + fake.reset_pre_plan_for_batch(batch) + assert batch.step is None and batch.preplanned_rids is None and batch.preplan_event is None + assert cleared == [True] From eef4b853ef8cb11a1d6938349b3c9131cc7128d4 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 205/219] kv: a snapshot of a stream reset since publishes the stream as it stands --- mstar/engine/resources/kv/manager.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/manager.py b/mstar/engine/resources/kv/manager.py index 1204b7513d..bfd2ac00e5 100644 --- a/mstar/engine/resources/kv/manager.py +++ b/mstar/engine/resources/kv/manager.py @@ -1905,7 +1905,11 @@ def publish_from_snapshot( ): """``publish`` as of the snapshot: the lengths are the snapshot's, so a step that committed in between (the next one is already running when - the worker finishes a publication) does not leak into it.""" + the worker finishes a publication) does not leak into it. A commit + only appends pages, so the snapshot's page count still names the + pages it saw. A reset or a fork in between replaces the list (the + stream's ``reset_generation`` moves); the snapshot describes pages + that no longer exist then, and the stream is published as it stands.""" del node_name, graph_walk if not snapshot: return None @@ -1918,6 +1922,10 @@ def publish_from_snapshot( stream = streams.get(label) if stream is None: continue + if stream.reset_generation != reset_generation: + seq_len = stream.stored_len + num_pages = len(stream.page_indices) + reset_generation = stream.reset_generation pages = list(stream.page_indices[:num_pages]) seq_info[label] = KVSequenceInfo( seq_len=seq_len, From 57ebad9c7ce0429cef1a4214dbc7851322bfcd47 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 206/219] test: publish from a snapshot after a reset --- test/modular/test_lazy_publish.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/modular/test_lazy_publish.py b/test/modular/test_lazy_publish.py index 949d0147d3..ec315f4ef6 100644 --- a/test/modular/test_lazy_publish.py +++ b/test/modular/test_lazy_publish.py @@ -243,3 +243,21 @@ def test_eager_finalize_still_publishes_every_request(monkeypatch): assert batch.per_request_info[1].resource_publish_info == {"kv": "kv-info:1"} assert batch.publish_snapshots is None eng.materialize_publish(batch, [1]) # nothing to finish, no error + + +def test_kv_snapshot_of_a_stream_reset_since_publishes_the_stream_as_it_stands(): + # a reset between the snapshot and its finish replaces the page list + # (the generation moves): the snapshot's page count would name pages + # of the old generation, so the stream is published as it is now + kv = _manager() + kv.ingest_request("r0", None) + _grow(kv, "r0", 40) + snap = kv.publish_snapshot_for_step("r0", "node", "w") + kv.reset_request("r0") + _grow(kv, "r0", 10) + live = kv.publish("r0", "node", "w") + lazy = kv.publish_from_snapshot("r0", snap, "node", "w") + assert _seq(lazy).seq_len == _seq(live).seq_len == 10 + assert _seq(lazy).page_indices == _seq(live).page_indices + assert _seq(lazy).reset_generation == _seq(live).reset_generation + assert _seq(lazy).latest_kv_transfer_info == _seq(live).latest_kv_transfer_info From b3e3b555b726eb3b8c97e9f18563d26b5d2b873c Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 207/219] worker: drop the spec yield hold, note why two slots suffice, note the snapshot contract --- mstar/worker/worker.py | 32 ++++++++++---------------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index ce8154cae2..96c149cb98 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -3133,6 +3133,11 @@ def _pp_stage(name: str) -> None: ) # A frame to a peer or the conductor carries this step's publication: # finish the snapshot for those rids (the stopped ones already are). + # The snapshot pins the lengths and page counts of step N, so the + # commit of N+1 that may have landed since does not leak in; only a + # reset or fork in between moves a stream past its snapshot, and the + # KV manager publishes such a stream as it stands (its generation + # says so). if info_rids: engine.materialize_publish(batch_N.node_batch, info_rids) for rid in info_rids: @@ -3692,16 +3697,14 @@ def run(self) -> None: spec_peek_for_fairness = ( os.environ.get("MSTAR_SPEC_PEEK_FOR_FAIRNESS", "1") == "1" ) - # How long the loop keeps speculating once another step is ready, so - # that at high concurrency the requests arriving over that window - # prefill in one step rather than one step each. 0 yields at once. - spec_yield_hold_s = float(os.environ.get("MSTAR_SPEC_YIELD_HOLD_MS", "0")) / 1000.0 - fairness_ready_since: list[float | None] = [None] consecutive_spec_steps = 0 yield_away_from_target: tuple[str, str] | None = None # Build N+2 while N+1 runs, between the two halves of N's # post-processing: its pre-plan then overlaps the routing half rather - # than the next launch. TP1 nodes only; a parallel node keeps the head + # than the next launch. Two slots (graph buffers, resource staging) + # still suffice: N's GPU work has completed by then (its outputs are + # what the post-processing reads), so N+2 takes N's slot while N+1 + # holds the other. TP1 nodes only; a parallel node keeps the head # broadcast in the order the followers settle it. early_spec = os.environ.get("MSTAR_EARLY_SPEC", "1") == "1" next_speculation: Speculation | None = None @@ -3736,22 +3739,7 @@ def _build_speculation( (pending.node_name, pending.graph_walk), ) ) - if other_ready: - now = _time.perf_counter() - if fairness_ready_since[0] is None: - fairness_ready_since[0] = now - must_yield_for_fairness = ( - now - fairness_ready_since[0] >= spec_yield_hold_s - ) - else: - fairness_ready_since[0] = None - must_yield_for_fairness = False - must_yield_away = ( - consecutive >= max_consecutive_spec - or must_yield_for_fairness - ) - if must_yield_away: - fairness_ready_since[0] = None + must_yield_away = consecutive >= max_consecutive_spec or other_ready if not must_yield_away and not admit_refused: if self.enable_nvtx: range_push("worker.speculate", synchronize=False) From 734972e33196914f53d305a689443b7baee744aa Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 208/219] docs: MSTAR_SPEC_YIELD_HOLD_MS removed, MSTAR_DECODE_BUCKETS added --- docs/environment_variables.rst | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index e3a821e4fc..bc96b8186b 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -295,13 +295,15 @@ Worker scheduling - Pre-plan the speculative batch's attention on a dedicated thread while the previous replay runs. ``0`` plans inline on the GPU thread. - * - ``MSTAR_QWEN35_DECODE_BUCKETS`` + * - ``MSTAR_DECODE_BUCKETS`` - ``1,2,4,8,16,32,48,64,96,128`` - - The decode batch sizes Qwen3.5 captures as CUDA graphs. A decode step + - The decode batch sizes a model captures as CUDA graphs (Qwen3.5 reads + it; a model with its own ladder reads it the same way). A decode step runs at the smallest bucket that holds its rows, so the ladder bounds both the rows per step and the graph memory; narrow it (for example ``1,2,4,8,16,32``) on a memory-tight GPU or for an A/B against the - former ladder. + former ladder. ``MSTAR_QWEN35_DECODE_BUCKETS`` overrides it for + Qwen3.5 alone. * - ``MSTAR_EARLY_SPEC`` - ``1`` - Build and pre-plan step N+2 while N+1 runs, between the stop and @@ -323,15 +325,6 @@ Worker scheduling - ``1024`` - Cap on back-to-back speculative steps before the leader yields to other ready work. - * - ``MSTAR_SPEC_YIELD_HOLD_MS`` - - ``0`` - - How long (ms) the speculative decode loop keeps going after another - step becomes ready, typically the prefill of requests that just - arrived. With ``0`` the loop yields at the next iteration and every - arrival tends to get a prefill step of its own; a hold of a few - milliseconds lets the arrivals of that window prefill together, which - is worth decode throughput at high concurrency and costs up to that - much time to first token. * - ``MSTAR_RUST_SEND_TIMING`` - unset - ``1`` makes the Rust runtime print, every 500 calls, how long the From 5b54fc77b9887b9571b19a553138625dfd901538 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 209/219] api: APIServerMessageBody names the message body union --- mstar/api_server/request_types.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mstar/api_server/request_types.py b/mstar/api_server/request_types.py index ba071ab3e9..6c5cb5da0c 100644 --- a/mstar/api_server/request_types.py +++ b/mstar/api_server/request_types.py @@ -66,12 +66,16 @@ class RequestFailed: status: int = 500 +# what an ``APIServerMessage`` carries; None for ``setup_done`` +APIServerMessageBody = ResultTensors | ResultTokens | RequestComplete | RequestFailed | None + + @dataclass class APIServerMessage: """Envelope for messages received by the API server.""" # "result_tensors" | "result_tokens" | "request_complete" | "request_failed" | "setup_done" message_type: str - body: ResultTensors | ResultTokens | RequestComplete | RequestFailed | None = None # None for setup_done + body: APIServerMessageBody = None @dataclass From 58381683818a69bfda4b3b7b539a35f651c5ce0e Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 210/219] tensor store: ColumnarTensorInfo.add_row --- mstar/communication/tensor_store.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/mstar/communication/tensor_store.py b/mstar/communication/tensor_store.py index f651983b98..e2699149e7 100644 --- a/mstar/communication/tensor_store.py +++ b/mstar/communication/tensor_store.py @@ -77,6 +77,20 @@ def dtype_index(self, dtype: torch.dtype) -> int: self.dtype_names.append(_dtype_name(dtype)) return idx + def add_row( + self, uuid: int, address: int, nbytes: int, ndims: int, + dims, stride, dtype_idx: int, + ) -> None: + """One tensor whose layout is already known (a row of a template): + its entry in every column, in the layout's order.""" + self.uuids.append(uuid) + self.addresses.append(address) + self.nbytes.append(nbytes) + self.ndims.append(ndims) + self.dims_flat.extend(dims) + self.stride_flat.extend(stride) + self.dtype_idx.append(dtype_idx) + def add_tensors_canonical( self, uuids: list[int], canonicals: list[torch.Tensor], From 37f96131039b73c25e4539383c56d8cd34f39e81 Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:23 -0700 Subject: [PATCH 211/219] tensors: the row fast path fills the columns through add_row --- mstar/communication/tensors.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index d9229a88d3..c3b1d31431 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -601,13 +601,7 @@ def store_row_outputs_batch( cpus: list[torch.Tensor | None] | None = [] if want_cpu else None sharded_uuids: list[int] = [] mint = self._uuid_minter.mint - add_uuid = columns.uuids.append - add_addr = columns.addresses.append - add_nbytes = columns.nbytes.append - add_ndims = columns.ndims.append - add_dims = columns.dims_flat.extend - add_stride = columns.stride_flat.extend - add_dtype = columns.dtype_idx.append + add_row = columns.add_row try: for rid in request_ids: i = row_of.get(rid) @@ -658,13 +652,7 @@ def store_row_outputs_batch( signal_idxs.append(s) num_tensors.append(1) canonicals.append(tensor) - add_uuid(uuid) - add_addr(base + i * step) - add_nbytes(nbytes) - add_ndims(ndims) - add_dims(dims) - add_stride(stride) - add_dtype(dtype_idx) + add_row(uuid, base + i * step, nbytes, ndims, dims, stride, dtype_idx) if want_cpu: cpus.append(host[i] if host is not None else None) if cfg is not None: From 3a814eb8a16f0908aba28ec7a99ad7695caa829a Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:47:24 -0700 Subject: [PATCH 212/219] qwen3_5: decode buckets from MSTAR_DECODE_BUCKETS, the model knob overrides --- mstar/model/qwen3_5/submodules.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 7b2dbaf006..5f144205b2 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -84,11 +84,13 @@ class LLMSubmodule(ARNodeSubmodule): # cap on rows per decode step: with 32 as the top, 64 requests in flight # ran as two alternating batches of 32 and each paid the whole per-step # host cost. A step above 32 rows pads to the next bucket. - # MSTAR_QWEN35_DECODE_BUCKETS="1,2,4,8,16,32" narrows the ladder (less graph - # memory, or the pre-#258 ladder for an A/B). + # MSTAR_DECODE_BUCKETS="1,2,4,8,16,32" narrows the ladder for every model + # that reads it (less graph memory, or the pre-#258 ladder for an A/B); + # MSTAR_QWEN35_DECODE_BUCKETS overrides it for this model alone. DECODE_CAPTURE_BATCH_SIZES = [ - int(x) for x in os.environ.get( - "MSTAR_QWEN35_DECODE_BUCKETS", "1,2,4,8,16,32,48,64,96,128", + int(x) for x in ( + os.environ.get("MSTAR_QWEN35_DECODE_BUCKETS") + or os.environ.get("MSTAR_DECODE_BUCKETS", "1,2,4,8,16,32,48,64,96,128") ).split(",") ] From d2137a5f3058932330cebf95e6b69fa958de890d Mon Sep 17 00:00:00 2001 From: merceod Date: Mon, 12 Oct 2026 08:54:57 -0700 Subject: [PATCH 213/219] test: the fairness settings without the hold --- test/modular/test_tp_early_spec_decision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_tp_early_spec_decision.py b/test/modular/test_tp_early_spec_decision.py index ec10f068d4..d1b7dc405c 100644 --- a/test/modular/test_tp_early_spec_decision.py +++ b/test/modular/test_tp_early_spec_decision.py @@ -51,7 +51,7 @@ def _pending(seq=3): def _fair(**kw): - return _SpecFairness(peek=True, hold_s=0.0, max_consecutive=1024, **kw) + return _SpecFairness(peek=True, max_consecutive=1024, **kw) def test_an_early_build_with_no_head_sends_nothing_and_leaves_the_decision(): From 41bc8e2e7b38d23f874c2904cf20b868235ce564 Mon Sep 17 00:00:00 2001 From: Stephen Duan Date: Tue, 13 Oct 2026 04:10:22 +0000 Subject: [PATCH 214/219] race on resizing of master buffer could cause loss of one batch of decode tokens --- mstar/engine/resources/sampler/utils.py | 30 ++++++++++++++++++++----- test/modular/test_sampler_last_token.py | 22 ++++++++++++++++++ 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index ccfb2f8d67..a492e63ca5 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1338,6 +1338,8 @@ class SamplerBuffers: # Master cache capacity (grown by doubling when more requests are # concurrently registered than the per-step buffer holds). _master_capacity: int = field(default=0, repr=False) + # Rows of the device masters; lags ``_master_capacity`` until the gpu thread grows them. + _device_capacity: int = field(default=0, repr=False) # the in-graph scatter (knob): padding rows address ``_trash_slot``, one # master row past the capacity that no request is ever given ingraph_scatter: bool = False @@ -1454,6 +1456,7 @@ def mk(dtype: torch.dtype, default: float) -> HostBuffer: seen_tokens=seen_tokens, min_p=mk(torch.float32, 0.0) if enable_min_p else None, _master_capacity=cap, + _device_capacity=cap, ingraph_scatter=ingraph_scatter, _trash_slot=cap if ingraph_scatter else 0, cg_slots=cg_slots, @@ -1521,6 +1524,8 @@ def _grow_master(self, new_capacity: int) -> None: Triggered when concurrently-registered requests exceed the current master capacity. Per-step buffers (sized to the cuda-graph max_bs) are NOT resized — the gather only reads ``padded_bs`` rows from master. + + Host rows only; the device masters grow on the gpu thread (``_sync_device_capacity``). """ if self.ingraph_scatter: raise RuntimeError( @@ -1530,12 +1535,23 @@ def _grow_master(self, new_capacity: int) -> None: ) for buf in self._scalar_buffers(): buf.grow_master(new_capacity) - self.offset.grow_master(new_capacity) - self.last_token.grow_master(new_capacity) - if self.seen_tokens is not None: - self.seen_tokens.grow_master(new_capacity) - self._free_slots.extend(range(self._master_capacity, new_capacity)) + old_capacity = self._master_capacity self._master_capacity = new_capacity + self._free_slots.extend(range(old_capacity, new_capacity)) + + def _sync_device_capacity(self) -> None: + """Grow the device masters on the gpu thread, behind its own scatters. + + Swapping them on the main thread let an in-flight scatter land in the old tensor. + """ + cap = self._master_capacity + if self._device_capacity >= cap: + return + self.offset.grow_master(cap) + self.last_token.grow_master(cap) + if self.seen_tokens is not None: + self.seen_tokens.grow_master(cap) + self._device_capacity = cap def register_request( self, rid: str, sampling_config: SamplingConfig | None = None, @@ -1600,6 +1616,7 @@ def stage_seen_token_masks( """ if self.seen_tokens is None: return + self._sync_device_capacity() # Batched for the same reason as sync_seen_token_masks: bs launches # become one. _init_slot still has to run per slot (it is CPU-side # bookkeeping plus a master-row write) before anything is staged. @@ -1720,6 +1737,7 @@ def gather_dynamic( The seen-token mask is large ([bs, V] bool); only gathered when the caller's graph actually applies the penalty in-graph (the Talker).""" + self._sync_device_capacity() if self._staged_rids.get(cg_slot) != (tuple(request_ids), padded_bs): self._stage_slot_idx(request_ids, padded_bs, cg_slot) with self._offset_reset_lock: @@ -1794,6 +1812,7 @@ def write_last_tokens(self, request_ids: list[str], tokens: torch.Tensor) -> Non requests' master rows. Every request must be registered.""" if not request_ids: return + self._sync_device_capacity() self.last_token.master.index_copy_( 0, self._slots_of(request_ids), tokens.to(self.last_token.master.dtype) ) @@ -1801,6 +1820,7 @@ def write_last_tokens(self, request_ids: list[str], tokens: torch.Tensor) -> Non def last_tokens_for(self, request_ids: list[str]) -> torch.Tensor: """``[len(request_ids)]`` last tokens of registered requests, for an eager step that has no staged slot row.""" + self._sync_device_capacity() return self.last_token.master.index_select(0, self._slots_of(request_ids)) def has_slot(self, rid: str) -> bool: diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py index 1df7393116..ba06227206 100644 --- a/test/modular/test_sampler_last_token.py +++ b/test/modular/test_sampler_last_token.py @@ -62,6 +62,27 @@ def test_growth_and_slot_reuse_keep_the_tokens(): assert bufs.last_tokens_for(["d", "b", "c"]).tolist() == [8, 6, 7] +def test_registration_growth_leaves_the_device_masters_to_the_gpu_thread(): + """A scatter that read the master before a registration grew it must still land.""" + bufs = _buffers(max_bs=2, cg_slots=1) + bufs.register_request("a") + bufs.register_request("b") + bufs.gather_static(["a", "b"], 2, 0) + bufs.gather_dynamic(["a", "b"], 2, 0) + bufs.slice_for_bs(2, 0)["last_token_buf"].copy_(torch.tensor([5, 6])) + # commit's scatter, split where the GIL can switch: it reads the master... + master = bufs.last_token.master + idx, real_bs = bufs._scatter_idx(0) + bufs.register_request("c") # ...the main thread grows past the capacity of 2... + assert bufs.last_token.master is master + master.index_copy_(0, idx, bufs.last_token.slot_view(0, real_bs)) # ...then writes + # the next step grows the device rows on the gpu thread, keeping the tokens + bufs.gather_static(["a", "b"], 2, 0) + bufs.gather_dynamic(["a", "b"], 2, 0) + assert bufs.last_token.master.shape[0] >= 3 and bufs.offset.master.shape[0] >= 3 + assert bufs.gather_last_tokens(0, 2).tolist() == [5, 6] + + def _ingraph(max_bs=4, cg_slots=2, slots=8): return SamplerBuffers.allocate( max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, @@ -108,4 +129,5 @@ def test_ingraph_scatter_masters_never_grow(): plain = SamplerBuffers.allocate(max_batch_size=2, device=torch.device("cpu")) for rid in ("a", "b", "c"): plain.register_request(rid) + plain.write_last_tokens(["c"], torch.tensor([7])) # a gpu-thread access grows the device rows assert plain.last_token.master.shape[0] >= 3 and not plain.ingraph_scatter From 2aaff1ac9d21bef7076e710abf76d0c73eb0dc99 Mon Sep 17 00:00:00 2001 From: Stephen Duan Date: Tue, 13 Oct 2026 04:36:51 +0000 Subject: [PATCH 215/219] python graph would emit no device-loopback decode tokens to client because did not support inline emit --- mstar/worker/worker.py | 3 ++- test/modular/test_worker_device_loopback.py | 9 ++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 9067988e0e..806234c147 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -3274,7 +3274,8 @@ def _all_on_device(self, engine, node_name: str, graph_walk: str, signals) -> bo """Whether every output signal of this step is one the node keeps on the device (cached per node and walk: it depends on the node's resources, not on the step).""" - if not signals: + # without inline emit, a rows-only step has no path for the client's tokens + if not signals or not self._graph_runtime.supports_inline_emit: return False key = (node_name, graph_walk) on_device = self._device_loopback.get(key) diff --git a/test/modular/test_worker_device_loopback.py b/test/modular/test_worker_device_loopback.py index 6165b3c966..b362954dca 100644 --- a/test/modular/test_worker_device_loopback.py +++ b/test/modular/test_worker_device_loopback.py @@ -17,9 +17,10 @@ def device_loopback_signals(self, node, walk): return self.on_device.get((node, walk), frozenset()) -def _worker(): +def _worker(inline_emit=True): w = Worker.__new__(Worker) w._device_loopback = {} + w._graph_runtime = SimpleNamespace(supports_inline_emit=inline_emit) return w @@ -36,6 +37,12 @@ def test_all_on_device_is_per_step_signals_and_cached_per_node_and_walk(): assert engine.calls == 2 +def test_no_rows_only_step_on_a_runtime_without_inline_emit(): + """Rows-only steps send the client's tokens inline; without it they were dropped.""" + engine = _Engine({("LLM", "decode"): frozenset({"text_inputs"})}) + assert not _worker(inline_emit=False)._all_on_device(engine, "LLM", "decode", ["text_inputs"]) + + def test_the_zero_tensor_store_shape_matches_the_route_layout(): rids, signals = [3, 5, 8], ["text_inputs"] stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) From d5706e51092c6a32c128b5293e31f4bff16acf51 Mon Sep 17 00:00:00 2001 From: Stephen Duan Date: Tue, 13 Oct 2026 19:08:25 +0000 Subject: [PATCH 216/219] lock to prevent race on host-side tensors that could lose request configs on submission --- mstar/engine/resources/sampler/utils.py | 27 ++++++++++++------- test/modular/test_sampler_last_token.py | 36 ++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index a492e63ca5..4dde8c08a4 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1376,6 +1376,8 @@ class SamplerBuffers: # thread, default stream) drains it. _pending_offset_reset: set[int] = field(default_factory=set, repr=False) _offset_reset_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + # Host master-row writes vs their regrow; a plan-thread write between copy and swap was lost. + _host_master_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) # Last-known config per rid — change-detect for ``update_request_config`` # so steady-state per-step calls do zero GPU work (for the scalar rows). _cached_config: dict[str, SamplingConfig] = field(default_factory=dict, repr=False) @@ -1509,13 +1511,16 @@ def _write_master_row(self, slot: int, cfg: SamplingConfig) -> None: # argmax (first index on ties), which from_probs returns for any # seed. See ``sample_cuda_graphable_gpu``. t, k, p = 0.0, 0, 1.0 - self.temperature.write_master_row(slot, t) - self.top_k.write_master_row(slot, k) - self.top_p.write_master_row(slot, p) - self.seed.write_master_row(slot, cfg.seed) - self.rep_penalty.write_master_row(slot, float(cfg.repetition_penalty)) - if self.min_p is not None: - self.min_p.write_master_row(slot, float(cfg.min_p) if cfg.temperature > 0 else 0.0) + with self._host_master_lock: + self.temperature.write_master_row(slot, t) + self.top_k.write_master_row(slot, k) + self.top_p.write_master_row(slot, p) + self.seed.write_master_row(slot, cfg.seed) + self.rep_penalty.write_master_row(slot, float(cfg.repetition_penalty)) + if self.min_p is not None: + self.min_p.write_master_row( + slot, float(cfg.min_p) if cfg.temperature > 0 else 0.0 + ) self._config_version += 1 def _grow_master(self, new_capacity: int) -> None: @@ -1525,7 +1530,8 @@ def _grow_master(self, new_capacity: int) -> None: master capacity. Per-step buffers (sized to the cuda-graph max_bs) are NOT resized — the gather only reads ``padded_bs`` rows from master. - Host rows only; the device masters grow on the gpu thread (``_sync_device_capacity``). + Host rows only, under the lock ``_write_master_row`` takes, so a plan-thread + ``_init_slot`` cannot land in the old array; device masters grow on the gpu thread. """ if self.ingraph_scatter: raise RuntimeError( @@ -1533,8 +1539,9 @@ def _grow_master(self, new_capacity: int) -> None: "sampler with the in-graph scatter on; raise MSTAR_SAMPLER_SLOTS " "or set MSTAR_SAMPLER_INGRAPH_SCATTER=0" ) - for buf in self._scalar_buffers(): - buf.grow_master(new_capacity) + with self._host_master_lock: + for buf in self._scalar_buffers(): + buf.grow_master(new_capacity) old_capacity = self._master_capacity self._master_capacity = new_capacity self._free_slots.extend(range(old_capacity, new_capacity)) diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py index ba06227206..b3a6b85d30 100644 --- a/test/modular/test_sampler_last_token.py +++ b/test/modular/test_sampler_last_token.py @@ -2,10 +2,13 @@ captured sample writes a per-step row, commit scatters it to the slot master, the next step gathers its input ids by slot; an eager step writes and reads the master by request.""" +import threading + +import numpy as np import pytest import torch -from mstar.engine.resources.sampler.utils import SamplerBuffers +from mstar.engine.resources.sampler.utils import HostBuffer, SamplerBuffers, SamplingConfig def _buffers(max_bs=4, cg_slots=2): @@ -83,6 +86,37 @@ def test_registration_growth_leaves_the_device_masters_to_the_gpu_thread(): assert bufs.gather_last_tokens(0, 2).tolist() == [5, 6] +def test_plan_thread_row_init_survives_a_concurrent_host_regrow(monkeypatch): + """A plan-thread ``_init_slot`` between the regrow's copy and swap must still land.""" + bufs = _buffers(max_bs=2, cg_slots=1) + copied, written = threading.Event(), threading.Event() + paused = [] + + def grow_master(self, new_capacity): + new = np.full(new_capacity, self.default, dtype=self.master.dtype) + new[: self.master.shape[0]] = self.master + if not paused: + paused.append(self) + copied.set() + written.wait(timeout=0.3) # times out once the write waits on the lock + self.master = new + + monkeypatch.setattr(HostBuffer, "grow_master", grow_master) + bufs.register_request("a", SamplingConfig(temperature=0.0)) + bufs.register_request("b") + + def plan(): + copied.wait() + bufs.gather_static(["a"], 1, 0) # inits a's row on the plan thread + written.set() + + t = threading.Thread(target=plan) + t.start() + bufs.register_request("c") # grows the host masters on the main thread + t.join() + assert bufs.temperature.master[bufs._rid_to_slot["a"]] == 0.0 + + def _ingraph(max_bs=4, cg_slots=2, slots=8): return SamplerBuffers.allocate( max_batch_size=max_bs, device=torch.device("cpu"), cg_slots=cg_slots, From b48904d6d2a6e42f8756775a491ae96efa7d2876 Mon Sep 17 00:00:00 2001 From: NSagan271 Date: Wed, 14 Oct 2026 06:40:51 +0000 Subject: [PATCH 217/219] Review comments: remove env vars for MSTAR_GC_FREEZE (always on), and MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT (logic removed) ; add a batch_postprocess function (currently unused) for added flexibility with rows_only ; fail request that doesn't fit into cuda graph scatter buffers instead of raising ; dummy rows map to dummy slot in cuda graph scatter buffers ; set rows_only to False if can_batch is False and there is no lease ; fix prepare_inputs for PD disaggregation with rows_only --- docs/environment_variables.rst | 18 ----- mstar/engine/engine.py | 50 ++++++++++-- mstar/engine/resources/runner.py | 22 ++++-- mstar/engine/resources/sampler/utils.py | 8 +- mstar/model/qwen3_5/submodules.py | 27 +++++-- mstar/model/submodule_base.py | 24 ++++-- mstar/utils/gc_freeze.py | 6 +- mstar/utils/knobs.py | 11 --- mstar/worker/worker.py | 73 +++++++----------- test/cpu-core.txt | 2 + test/modular/test_engine_rows_only_tail.py | 81 ++++++++++++++++++++ test/modular/test_engine_uniform_inputs.py | 30 +++++++- test/modular/test_gc_freeze.py | 5 +- test/modular/test_ingest_rollback.py | 75 ++++++++++++++++++ test/modular/test_qwen3_5_device_loopback.py | 63 +++++++++++---- test/modular/test_sampler_last_token.py | 14 ++++ test/modular/test_worker_launch_signal.py | 67 ---------------- 17 files changed, 390 insertions(+), 186 deletions(-) create mode 100644 test/modular/test_engine_rows_only_tail.py create mode 100644 test/modular/test_ingest_rollback.py delete mode 100644 test/modular/test_worker_launch_signal.py diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index fdfc18756d..892875c60a 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -329,15 +329,6 @@ Worker scheduling 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_GC_FREEZE`` - - ``1`` - - Once a process has finished its set-up (the API server and the - conductor when every worker is ready, the worker after warm-up and - graph capture), collect once and ``gc.freeze()`` the heap, so later - gen-2 collections do not walk the long-lived set-up objects. Measured - on the API server process: 200 ms per gen-2 pass every ~20 s under - load before, which stalled request admission and token delivery for - that long. ``0`` leaves the collector alone. * - ``MSTAR_DEVICE_LOOPBACK`` - ``1`` - Whether a decode node may keep its loop-back token on the device: @@ -357,15 +348,6 @@ Worker scheduling 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_LAUNCH_SIGNAL_AFTER_COMMIT`` - - ``0`` - - Release the thread that submitted a step only once the gpu thread - has also committed it and staged its outputs, instead of at the - replay launch. Every torch call after the launch releases the GIL and - the woken main thread keeps it for its whole postprocess stretch, so - the first such call waits that long; signalling later costs the main - thread the uncontended commit and collect instead. 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 diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 1e2a465047..9e0d578119 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -503,6 +503,9 @@ def load_model( for node in relevant_nodes: node_to_resources.setdefault(node, []).append(spec.resource_key) + # 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 + self._loopback_steps: dict[int, tuple[str, str]] = {} self._runner = StepRunner( self._resources, # every node, including one that owns nothing (Code2Wav) — an @@ -760,8 +763,11 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: # read leniently, as the probe itself was only reached per prepared row keyed = getattr(self, "_keyed_walks", None) probe_prefix = bool(keyed) and walk in keyed.get(batch.node_name, ()) + loopback = self._loopback_rows(batch) if not probe_prefix: - uniform = submodule.uniform_row_inputs(walk) + uniform = submodule.uniform_row_inputs( + walk, has_loopback=loopback is not None and all(loopback), + ) if uniform is not None: # one object for all the step's rows: the per-row call would # return an equal one each time (128 calls at a full decode @@ -772,7 +778,7 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: batch=batch, model_inputs=node_inputs ) return - for rid in batch.request_ids: + for i, rid in enumerate(batch.request_ids): try: req_inputs = prepare( graph_walk=walk, @@ -784,6 +790,7 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: # look-ahead frames a token encoder withholds) is_final_stream_chunk=rid in final_rids, input_metadata=metadata.get(rid, EMPTY_INPUT_METADATA), + has_loopback=loopback is not None and loopback[i], ) if req_inputs is not None and probe_prefix: req_inputs = self._skip_cached_prefix(batch, rid, req_inputs) @@ -805,6 +812,26 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None: batch=batch, model_inputs=node_inputs ) + def _loopback_rows(self, batch: ExecutingBatch) -> list[bool] | None: + """Per row, whether its input token is on the device (see + ``_loopback_steps``); None when no row's is.""" + steps = self._loopback_steps + if not steps: + return None + key = (batch.node_name, batch.step_context.graph_walk) + get = steps.get + return [get(rid) == key for rid in batch.request_ids] + + def _record_loopback(self, batch: ExecutingBatch) -> None: + steps = self._loopback_steps + if batch.rows_only: + steps.update(dict.fromkeys( + batch.request_ids, (batch.node_name, batch.step_context.graph_walk), + )) + elif steps: + for rid in batch.request_ids: + steps.pop(rid, None) + def inline_client_signals( self, node_name: str, graph_walk: str, ) -> dict[str, str]: @@ -914,9 +941,13 @@ def exec( # in order; an unbatchable walk with >1 request can't, so each # request runs its own full cycle and merges. if not batch.running_batched and len(batch.request_ids) > 1 and lease is None: + # per-request outputs: the step cannot stay on the rows + batch.rows_only = False batch.outputs = self._exec_per_request(batch) else: batch.outputs = self._exec_single(batch) + if batch.admit_error is None: + self._record_loopback(batch) batch.outputs_ready.set() return batch.outputs except Exception: @@ -1411,9 +1442,17 @@ def exec_and_postprocess( outputs = self.exec(batch) # A failed admit means no forward ran and every rid's outputs are empty, # so the tail has nothing to consume; the worker re-drives the step. - # A rows-only step has no per-request outputs for a tail to work on. - if batch.admit_error is None and not batch.rows_only: - self.postprocess_batch(batch, outputs) + # A rows-only step has no per-request outputs; only a batch tail runs. + if batch.admit_error is None: + if not batch.rows_only: + self.postprocess_batch(batch, outputs) + else: + submodule = self._submodules[batch.node_name].submodule + # no submodule overrides it as of now + if type(submodule).batch_postprocess is not NodeSubmodule.batch_postprocess: + submodule.batch_postprocess( + outputs, batch.per_request_info_wrapped, batch.inputs, + ) return outputs def check_stop_for_batch( @@ -2213,6 +2252,7 @@ def add_request( self._runner.ingest_request(request_id, overrides) def remove_request(self, request_id: str) -> None: + self._loopback_steps.pop(request_id, None) self._runner.remove_request(request_id) for submodule_mgmt in self._submodules.values(): submodule_mgmt.submodule.cleanup_request(request_id) diff --git a/mstar/engine/resources/runner.py b/mstar/engine/resources/runner.py index e0f606e22c..56787ad5bf 100644 --- a/mstar/engine/resources/runner.py +++ b/mstar/engine/resources/runner.py @@ -263,11 +263,23 @@ def ingest_request( self, rid: str, overrides: Mapping[str, ResourceReqConfig] | None = None, ) -> None: - """open state on all resources; `overrides` are per-resource""" - for key in self._order: - self._resources[key].ingest_request( - rid, None if overrides is None else overrides.get(key) - ) + """open state on all resources; `overrides` are per-resource. All or + nothing: a resource that raises leaves the request on none of them.""" + for i, key in enumerate(self._order): + try: + self._resources[key].ingest_request( + rid, None if overrides is None else overrides.get(key) + ) + except Exception: + # the raising resource may have taken part of it too + for done in self._order[:i + 1]: + try: + self._resources[done].remove_request(rid) + except Exception: + logger.exception( + "rolling back request %s on resource %s failed", rid, done, + ) + raise def remove_request(self, rid: str) -> None: for key in self._order: diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index 9b495c7088..c5206acae3 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -1581,6 +1581,9 @@ def register_request( self, rid: str, sampling_config: SamplingConfig | None = None, ) -> None: """Allocate a slot for ``rid`` and seed its master row.""" + # a capture padding handle (negative) takes no slot; it reads slot 0 + if isinstance(rid, int) and rid < 0: + return if rid in self._rid_to_slot: # Re-registration: just refresh the config in place. if sampling_config is not None: @@ -1844,7 +1847,10 @@ def gather_last_tokens(self, cg_slot: int, padded_bs: int) -> torch.Tensor: ) def _slots_of(self, request_ids: list[str]) -> torch.Tensor: - slots = [self._rid_to_slot[rid] for rid in request_ids] + slots = [ + 0 if isinstance(rid, int) and rid < 0 else self._rid_to_slot[rid] + for rid in request_ids + ] return torch.tensor(slots, dtype=torch.long, device=self.last_token.master.device) def write_last_tokens(self, request_ids: list[str], tokens: torch.Tensor) -> None: diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index 3c8d7290da..2acf405560 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -286,21 +286,24 @@ def prepare_inputs( graph_walk: str, fwd_info: CurrentForwardPassInfo, inputs: NameToTensorList, + has_loopback: bool = False, **kwargs: Any, ) -> ARNodeInputs: if graph_walk == "prefill_vision": return self._vision_inputs(fwd_info, inputs) - if graph_walk == "decode" and self._device_loopback(): + if graph_walk == "decode" and has_loopback and self._device_loopback(): # one token per row, read off the sampler's slot master in # `preprocess`; the routed signal is empty (or ignored) here return ARNodeInputs(input_seq_len=1) input_ids = inputs["text_inputs"][0] return ARNodeInputs(input_seq_len=input_ids.shape[0], input_ids=input_ids) - def uniform_row_inputs(self, graph_walk: str) -> ARNodeInputs | None: + def uniform_row_inputs( + self, graph_walk: str, has_loopback: bool = False, + ) -> ARNodeInputs | None: # the loop-back decode row above, built once for the step's rows - if graph_walk == "decode" and self._device_loopback(): + if graph_walk == "decode" and has_loopback and self._device_loopback(): return ARNodeInputs(input_seq_len=1) return None @@ -353,12 +356,26 @@ def preprocess( inputs: list[ARNodeInputs], ) -> dict[str, torch.Tensor | Any]: out: dict[str, torch.Tensor | Any] = {} - if inputs[0].input_ids is None and inputs[0].input_embeds is None: + first = inputs[0] + if first.input_embeds is None and ( + first.input_ids is None or any(inp.input_ids is None for inp in inputs) + ): # device loop-back: every row of a decode step is its request's # last sampled token, read by slot (padding rows read slot 0) sampler: SamplerResource = engine_inputs.resources[SAMPLER] t0 = perf_counter() if PHASE_PERIOD else 0.0 - out["input_ids"] = sampler.loopback_tokens(engine_inputs.request_ids) + ids = sampler.loopback_tokens(engine_inputs.request_ids) + # rows whose last step ran elsewhere brought their token + routed = [ + i for i, inp in enumerate(inputs) + if inp.input_ids is not None and inp.input_seq_len + ] + if routed: + ids.index_copy_( + 0, torch.tensor(routed, device=ids.device), + torch.cat([inputs[i].input_ids for i in routed]).to(ids.dtype), + ) + out["input_ids"] = ids if PHASE_PERIOD: phase_record("engine.preprocess.loopback", perf_counter() - t0) elif inputs[0].input_ids is not None: diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index 59044435ab..efbe85bbc5 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -696,14 +696,18 @@ def filter_batched_output( a fixed key set for graph compat, so the filtering happens here.""" return outputs - def uniform_row_inputs(self, graph_walk: str) -> "NodeInputs | None": + def uniform_row_inputs( + self, graph_walk: str, has_loopback: bool = False, + ) -> "NodeInputs | None": """The inputs every row of this walk gets, when ``prepare_inputs`` would return an equal object for each regardless of the request, its routed signals and its metadata (a decode walk that reads its token off the device). The engine then prepares the batch with one call and shares the object across the step's rows, so it must not be mutated - per row. None (the default) prepares row by row.""" - del graph_walk + per row. ``has_loopback``: every row's last step on this node kept its + loop-back signals on the device (``device_loopback_signals``). None + (the default) prepares row by row.""" + del graph_walk, has_loopback return None def step_is_input_free(self, graph_walk: str) -> bool: @@ -720,8 +724,9 @@ def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: keeps on the device itself (a sampler's slot master), so the worker routes them with no tensor: the next step reads the value in ``preprocess`` and every fallback batch arrives with the signal - empty. Only sound when ``prepare_inputs`` accepts the empty signal on - this walk. On a walk whose output signals all stay on the device the + empty. ``prepare_inputs`` is told per row (``has_loopback``) whether + its last step here kept the value on the device; a row without it + (e.g. its previous step ran on another rank) brings the tensor. On a walk whose output signals all stay on the device the engine builds no per-request output dicts and does not call ``postprocess``: the stop check, the inline emit and the next step's speculation read the step's rows instead. Empty by default; @@ -877,6 +882,15 @@ def postprocess( """ return + def batch_postprocess( + self, outputs: BatchedModelOutput, + request_infos: dict[int, CurrentForwardPassInfo], + inputs: list[NodeInputs], + ) -> None: + """``postprocess`` for a rows-only step, which has no per-request + outputs: works on the rows of ``outputs`` in place. Optional; only + called when overridden.""" + def check_stop( self, request_id: str, request_info: CurrentForwardPassInfo, diff --git a/mstar/utils/gc_freeze.py b/mstar/utils/gc_freeze.py index 062c9ce0b0..53be3c2c02 100644 --- a/mstar/utils/gc_freeze.py +++ b/mstar/utils/gc_freeze.py @@ -11,17 +11,13 @@ """ import gc import logging -import os logger = logging.getLogger(__name__) def freeze_after_setup(who: str) -> int: """``gc.collect()`` then ``gc.freeze()`` for a process that finished its - set-up. Returns the number of objects frozen; 0 and no change when - ``MSTAR_GC_FREEZE=0``.""" - if os.environ.get("MSTAR_GC_FREEZE", "1") != "1": - return 0 + set-up. Returns the number of objects frozen.""" gc.collect() gc.freeze() frozen = gc.get_freeze_count() diff --git a/mstar/utils/knobs.py b/mstar/utils/knobs.py index 7d6831d391..e6e0df4388 100644 --- a/mstar/utils/knobs.py +++ b/mstar/utils/knobs.py @@ -10,17 +10,6 @@ def device_loopback_enabled() -> bool: return os.environ.get("MSTAR_DEVICE_LOOPBACK", "1") == "1" -def launch_signal_after_commit() -> bool: - """``MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT`` (default 0): release the thread - that submitted a step only once the gpu thread has also committed the - step and staged its outputs, instead of at the replay launch. Every torch - call after the launch releases the GIL, and the woken main thread keeps - it for its whole postprocess stretch, so the first such call waits that - long; signalling later costs the main thread the uncontended commit and - collect (a fraction of a millisecond) instead.""" - return os.environ.get("MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT", "0") == "1" - - def sampler_ingraph_scatter() -> bool: """``MSTAR_SAMPLER_INGRAPH_SCATTER`` (default 0): the captured sampler gathers its RNG offsets from the slot masters and scatters the advanced diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 24708893d5..1af9ab1e45 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -73,7 +73,6 @@ WorkerMessage, WorkerMessageType, ) -from mstar.utils.knobs import launch_signal_after_commit from mstar.utils.profiler import PHASE_PERIOD, nvtx_enabled, phase_buffer, range_pop, range_push from mstar.worker.engine_manager import EngineManager from mstar.worker.micro_scheduler import MicroScheduler, ScheduledBatch @@ -277,8 +276,6 @@ class Worker: # Empty for a text model, which lets a speculation build skip its three # per-row stream polls; None (a bare instance) means unknown, so poll. _stream_buffer_rids: set[int] | None = None - # a bare instance keeps the launch signal where the engine sets it - _launch_signal_after_commit: bool = False def __init__( self, @@ -308,7 +305,6 @@ def __init__( self._device_loopback: dict[tuple[str, str], frozenset[str]] = {} self._host_rows_slot = 0 self._stream_buffer_rids = set() - self._launch_signal_after_commit = launch_signal_after_commit() # Per-phase wall-clock timing (MSTAR_PHASE_TIMING). On the worker # rather than in run()'s scope because the GPU and plan threads @@ -690,12 +686,19 @@ def _add_new_request(self, body: NewRequest) -> None: self._last_active[(request_id, node_name)] = now self.request_state.add_request(request_id, body.request_info) - self.engine_manager.add_request( - request_id, body.request_info.resource_configs, - ) self.tensor_manager.register_request( request_id, self._graph_runtime.get_sharding_config(request_id), ) + try: + self.engine_manager.add_request( + request_id, body.request_info.resource_configs, + ) + except Exception as err: + # the engine rolled the ingest back; the conductor's REMOVE tears down the rest + logger.exception("Worker %s: ingest of request %s failed", self.worker_id, body.request_id) + self._fail_requests({request_id: f"Error in worker: {type(err).__name__}: {err}"}) + self._process_unprocessed(body.request_id) + return # Create StreamBuffers for consumer connections on this worker. A request # with several partitions here arrives once per partition: keep the @@ -735,10 +738,13 @@ def _add_new_request(self, body: NewRequest) -> None: self._graph_runtime.ingest_inputs_batch( self._one_rids_edges(request_id, signal_only), can_buffer=True, ) - # process messages that may have came in out-of-order - if body.request_id in self._unprocessed_messages: - self._process_message_list(self._unprocessed_messages[body.request_id]) - del self._unprocessed_messages[body.request_id] + self._process_unprocessed(body.request_id) + + def _process_unprocessed(self, wire_rid: str) -> None: + """Process messages for ``wire_rid`` that came in before its NEW_REQUEST.""" + if wire_rid in self._unprocessed_messages: + self._process_message_list(self._unprocessed_messages[wire_rid]) + del self._unprocessed_messages[wire_rid] def _remove_request(self, body: RemoveRequest) -> None: @@ -1899,12 +1905,18 @@ def _execute_on_gpu_thread( self._phase_record( f"rows.{batch.graph_walk}#", len(node_batch.request_ids), ) - outputs = self._exec_step(engine, node_batch, execution_stream) - if gpu_start is not None: - gpu_end = torch.Event(enable_timing=True) - gpu_end.record(execution_stream) - self._gpu_spans.append((gpu_start, gpu_end)) - self._drain_gpu_spans() + with self._span("worker.gpu_thread.exec"): + outputs = engine.exec_and_postprocess(node_batch) + if execution_stream is not None: + self._stage_host_rows(outputs) + event = torch.Event() + event.record(execution_stream) + node_batch.completion_event = event + if gpu_start is not None: + gpu_end = torch.Event(enable_timing=True) + gpu_end.record(execution_stream) + self._gpu_spans.append((gpu_start, gpu_end)) + self._drain_gpu_spans() return outputs finally: # Safety net: a step that raised before the forward would otherwise @@ -3250,33 +3262,6 @@ def _inline_emit_values( return None, [] return None, [] - def _exec_step(self, engine, node_batch: ExecutingBatch, execution_stream): - """The step on the gpu thread: the engine's forward and commit, the - staged host rows, the completion event. - - With ``MSTAR_LAUNCH_SIGNAL_AFTER_COMMIT`` the engine sees no event to - set at the replay launch; the submitter is released here afterwards, - once commit, collect and the staging are done, so none of those torch - calls has to win the GIL back from the thread it woke. - """ - deferred_signal = None - if self._launch_signal_after_commit and node_batch.launch_started_event is not None: - deferred_signal = node_batch.launch_started_event - node_batch.launch_started_event = None - try: - with self._span("worker.gpu_thread.exec"): - outputs = engine.exec_and_postprocess(node_batch) - if execution_stream is not None: - self._stage_host_rows(outputs) - event = torch.Event() - event.record(execution_stream) - node_batch.completion_event = event - finally: - if deferred_signal is not None: - node_batch.launch_started_event = deferred_signal - deferred_signal.set() - return outputs - def _stage_host_rows(self, outputs: "BatchedModelOutput") -> None: """Start the device-to-host copy of the step's row-addressed stop buffers here on the gpu thread, on the execution stream right behind diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 751ab4df26..e8816f0d7d 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -72,3 +72,5 @@ test/modular/test_fast_tokenize.py test/modular/test_wire_codec.py test/modular/test_wire_security.py test/modular/test_model_kwargs_tensor_transport.py +test/modular/test_engine_rows_only_tail.py +test/modular/test_ingest_rollback.py diff --git a/test/modular/test_engine_rows_only_tail.py b/test/modular/test_engine_rows_only_tail.py new file mode 100644 index 0000000000..adfb7961bc --- /dev/null +++ b/test/modular/test_engine_rows_only_tail.py @@ -0,0 +1,81 @@ +"""A rows-only step runs only a submodule's batch tail, and only when it has +one; a step that falls back to the per-request path is not rows-only.""" +import threading +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.engine import Engine # noqa: E402 +from mstar.model.submodule_base import BatchedModelOutput, NodeSubmodule # noqa: E402 + + +class _Plain(NodeSubmodule): + def prepare_inputs(self, *a, **kw): + return None + + def forward(self, *a, **kw): + return None + + +class _WithTail(_Plain): + def batch_postprocess(self, outputs, request_infos, inputs): + self.seen = (outputs, request_infos, inputs) + + +def _engine(sub): + eng = Engine.__new__(Engine) + eng._submodules = {"llm": SimpleNamespace(submodule=sub)} + eng._loopback_steps = {} + eng.per_request_tails = [] + eng.postprocess_batch = lambda batch, outputs: eng.per_request_tails.append(batch) + return eng + + +def _batch(rows_only): + return SimpleNamespace( + node_name="llm", admit_error=None, rows_only=rows_only, + per_request_info_wrapped={1: object()}, inputs=[object()], + ) + + +@pytest.mark.parametrize("sub_cls", [_Plain, _WithTail]) +def test_rows_only_runs_the_batch_tail_only_when_overridden(sub_cls): + sub = sub_cls.__new__(sub_cls) + eng = _engine(sub) + out = BatchedModelOutput() + eng.exec = lambda batch: out + batch = _batch(rows_only=True) + assert eng.exec_and_postprocess(batch) is out + assert eng.per_request_tails == [] + if sub_cls is _WithTail: + assert sub.seen == (out, batch.per_request_info_wrapped, batch.inputs) + eng.exec_and_postprocess(_batch(rows_only=False)) + assert len(eng.per_request_tails) == 1 + + +def test_the_per_request_fallback_clears_rows_only(): + eng = _engine(_Plain.__new__(_Plain)) + eng._enable_nvtx = False + eng._enable_profile = False + eng._device = SimpleNamespace(type="cpu") + eng._autocast_dtype_for = lambda sub: None + eng.reserve_replay_slot = lambda batch: None + seen = {} + + def _per_request(batch): + seen["rows_only"] = batch.rows_only + return BatchedModelOutput() + + eng._exec_per_request = _per_request + batch = SimpleNamespace( + node_name="llm", request_ids=[1, 2], rows_only=True, running_batched=False, + step_context=SimpleNamespace(slot_lease=None, graph_walk="decode"), + exec_timings=SimpleNamespace(start=None), admit_error=None, + outputs_ready=threading.Event(), release_waiters=lambda: None, + preplanned_rids=None, preplan_event=None, + ) + eng.exec(batch) + assert seen["rows_only"] is False and batch.rows_only is False + assert eng._loopback_steps == {} diff --git a/test/modular/test_engine_uniform_inputs.py b/test/modular/test_engine_uniform_inputs.py index 7adeca9247..9cc184caf1 100644 --- a/test/modular/test_engine_uniform_inputs.py +++ b/test/modular/test_engine_uniform_inputs.py @@ -39,12 +39,16 @@ class _Sub: def __init__(self, uniform=None): self.uniform = uniform self.calls = [] + self.row_loopback = [] + self.uniform_loopback = None - def uniform_row_inputs(self, graph_walk): + def uniform_row_inputs(self, graph_walk, has_loopback=False): + self.uniform_loopback = has_loopback return self.uniform def prepare_inputs(self, **kw): self.calls.append(kw["graph_walk"]) + self.row_loopback.append(kw["has_loopback"]) return object() def can_batch(self, batch, model_inputs): @@ -55,6 +59,7 @@ def _engine(sub): eng = Engine.__new__(Engine) eng._submodules = {"llm": SimpleNamespace(submodule=sub, resources={})} eng._enable_nvtx = False + eng._loopback_steps = {} return eng @@ -72,3 +77,26 @@ def test_other_walks_still_prepare_row_by_row(): batch = _Batch([1, 2]) _engine(sub).prepare_inputs(batch) assert sub.calls == ["decode", "decode"] and len(batch.inputs) == 2 + + +def test_loopback_is_per_row_and_tracks_each_request_s_last_step(): + sub = _Sub(uniform=None) + eng = _engine(sub) + ran = _Batch([1, 2]) + ran.rows_only = True + eng._record_loopback(ran) + eng.prepare_inputs(_Batch([1, 2, 3])) + # 3 never ran rows-only here (e.g. its prefill ran on another rank) + assert sub.uniform_loopback is False and sub.row_loopback == [True, True, False] + eng.prepare_inputs(_Batch([1, 2])) + assert sub.uniform_loopback is True + # a step that is not rows-only clears the record; so does another walk's + other = _Batch([2]) + other.rows_only = False + eng._record_loopback(other) + sub.row_loopback.clear() + eng.prepare_inputs(_Batch([1, 2])) + assert sub.row_loopback == [True, False] + sub.row_loopback.clear() + eng.prepare_inputs(_Batch([1], walk="prefill")) + assert sub.row_loopback == [False] diff --git a/test/modular/test_gc_freeze.py b/test/modular/test_gc_freeze.py index bdbb743bc5..ff4dce574d 100644 --- a/test/modular/test_gc_freeze.py +++ b/test/modular/test_gc_freeze.py @@ -5,11 +5,8 @@ from mstar.utils.gc_freeze import freeze_after_setup -def test_freeze_after_setup_moves_the_heap_and_honours_the_knob(monkeypatch): +def test_freeze_after_setup_moves_the_heap(): gc.unfreeze() - monkeypatch.setenv("MSTAR_GC_FREEZE", "0") - assert freeze_after_setup("test") == 0 and gc.get_freeze_count() == 0 - monkeypatch.setenv("MSTAR_GC_FREEZE", "1") keep = [[i] for i in range(100)] # live objects the collector tracks frozen = freeze_after_setup("test") try: diff --git a/test/modular/test_ingest_rollback.py b/test/modular/test_ingest_rollback.py new file mode 100644 index 0000000000..9fc68919b9 --- /dev/null +++ b/test/modular/test_ingest_rollback.py @@ -0,0 +1,75 @@ +"""A request whose ingest raises on one resource is left on none of them, and +the worker fails only that request instead of the batch in flight.""" +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from mstar.engine.resources.runner import StepRunner # noqa: E402 +from mstar.worker.worker import Worker # noqa: E402 + + +class _Res: + def __init__(self, name, log, fail=False): + self.name, self.log, self.fail = name, log, fail + + def ingest_request(self, rid, overrides=None): + self.log.append(("ingest", self.name, rid)) + if self.fail: + raise RuntimeError("over capacity") + + def remove_request(self, rid): + self.log.append(("remove", self.name, rid)) + + +def _runner(resources): + runner = StepRunner.__new__(StepRunner) + runner._resources = {r.name: r for r in resources} + runner._order = [r.name for r in resources] + return runner + + +def test_a_raising_ingest_rolls_back_every_resource_it_reached(): + log = [] + runner = _runner([_Res("kv", log), _Res("sampler", log, fail=True), _Res("pos", log)]) + with pytest.raises(RuntimeError, match="over capacity"): + runner.ingest_request(7) + assert log == [ + ("ingest", "kv", 7), ("ingest", "sampler", 7), + ("remove", "kv", 7), ("remove", "sampler", 7), + ] + + +def test_the_worker_fails_only_the_request_whose_ingest_raised(): + w = Worker.__new__(Worker) + w.worker_id = "w0" + w._draining_rids = set() + w._unprocessed_messages = {"wire": ["parked"]} + w._last_active = {} + processed, failed, registered = [], [], [] + w._graph_runtime = SimpleNamespace( + add_request=lambda **kw: 5, get_sharding_config=lambda rid: None, + ) + w.engine_manager = SimpleNamespace( + evictable_nodes=lambda: [], + add_request=lambda rid, cfg: (_ for _ in ()).throw(RuntimeError("over capacity")), + ) + w.request_state = SimpleNamespace(add_request=lambda rid, info: None) + w.tensor_manager = SimpleNamespace( + register_request=lambda rid, sharding: registered.append(rid), + start_read_tensors=lambda *a, **kw: pytest.fail("no reads for a failed request"), + ) + w._fail_requests = failed.append + w._process_message_list = processed.extend + body = SimpleNamespace( + request_id="wire", partition_worker_graph_ids=None, worker_graph_to_workers={}, + request_info=SimpleNamespace( + partition_name="p", graph_walk="g", resource_configs=None, rid_handle=None, + ), + ) + Worker._add_new_request(w, body) + assert registered == [5] and len(failed) == 1 and list(failed[0]) == [5] + assert "over capacity" in failed[0][5] + # messages parked before the NEW still go through (a parked REMOVE among them) + assert processed == ["parked"] and "wire" not in w._unprocessed_messages diff --git a/test/modular/test_qwen3_5_device_loopback.py b/test/modular/test_qwen3_5_device_loopback.py index d98b282eb9..89b5dcc3f5 100644 --- a/test/modular/test_qwen3_5_device_loopback.py +++ b/test/modular/test_qwen3_5_device_loopback.py @@ -27,8 +27,10 @@ def _sub(sampler): return sub -def _prepare(sub, walk, inputs): - return LLMSubmodule.prepare_inputs(sub, graph_walk=walk, fwd_info=None, inputs=inputs) +def _prepare(sub, walk, inputs, has_loopback=True): + return LLMSubmodule.prepare_inputs( + sub, graph_walk=walk, fwd_info=None, inputs=inputs, has_loopback=has_loopback, + ) def test_decode_takes_no_input_ids_when_the_sampler_keeps_them(): @@ -42,6 +44,13 @@ def test_decode_takes_no_input_ids_when_the_sampler_keeps_them(): assert LLMSubmodule.device_loopback_signals(sub, "prefill_text") == frozenset() +def test_a_row_without_loopback_reads_its_routed_token(): + # e.g. its prefill ran on another rank: the master does not hold its token + sub = _sub(_Sampler(masters=True)) + out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}, has_loopback=False) + assert out.input_ids.tolist() == [7] and out.input_seq_len == 1 + + def test_without_masters_or_with_the_knob_off_the_tensor_path_stays(monkeypatch): sub = _sub(_Sampler(masters=False)) out = _prepare(sub, "decode", {"text_inputs": [torch.tensor([7])]}) @@ -56,15 +65,7 @@ def test_without_masters_or_with_the_knob_off_the_tensor_path_stays(monkeypatch) assert LLMSubmodule.device_loopback_signals(sub, "decode") == frozenset() -def test_preprocess_reads_the_step_tokens_off_the_sampler(): - sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 0, 0])) - sub = _sub(sampler) - rows = [_prepare(sub, "decode", {}) for _ in range(4)] - engine_inputs = SimpleNamespace( - resources={SAMPLER: sampler}, request_ids=["r1", "r2", "pad1", "pad2"], - ) - seen = {} - +def _stub_model(sub, seen): def _position_ids_3d(inputs): seen["n"] = len(inputs) return torch.zeros(3, len(inputs)) @@ -74,6 +75,34 @@ def _position_ids_3d(inputs): embed_tokens=SimpleNamespace(weight=torch.zeros(1)), )) sub._position_ids_3d = _position_ids_3d + + +def test_preprocess_mixes_routed_rows_into_the_sampler_tokens(): + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 5, 0])) + sub = _sub(sampler) + rows = [ + _prepare(sub, "decode", {"text_inputs": [torch.tensor([9])]}, has_loopback=False), + _prepare(sub, "decode", {}), + _prepare(sub, "decode", {"text_inputs": [torch.tensor([8])]}, has_loopback=False), + _prepare(sub, "decode", {}), + ] + engine_inputs = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2", "r3", "pad1"], + ) + _stub_model(sub, {}) + out = LLMSubmodule.preprocess(sub, "decode", engine_inputs, rows) + assert out["input_ids"].tolist() == [9, 4, 8, 0] + + +def test_preprocess_reads_the_step_tokens_off_the_sampler(): + sampler = _Sampler(masters=True, tokens=torch.tensor([3, 4, 0, 0])) + sub = _sub(sampler) + rows = [_prepare(sub, "decode", {}) for _ in range(4)] + engine_inputs = SimpleNamespace( + resources={SAMPLER: sampler}, request_ids=["r1", "r2", "pad1", "pad2"], + ) + seen = {} + _stub_model(sub, seen) out = LLMSubmodule.preprocess(sub, "decode", engine_inputs, rows) assert out["input_ids"].tolist() == [3, 4, 0, 0] assert sampler.asked == ["r1", "r2", "pad1", "pad2"] @@ -84,10 +113,14 @@ def test_the_loop_back_decode_row_is_uniform(monkeypatch): from mstar.model.submodule_base import ARNodeInputs sub = _sub(_Sampler(masters=True)) - row = LLMSubmodule.uniform_row_inputs(sub, "decode") + row = LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=True) assert isinstance(row, ARNodeInputs) and row.input_seq_len == 1 assert row.input_ids is None and row.input_embeds is None - assert LLMSubmodule.uniform_row_inputs(sub, "prefill_text") is None - assert LLMSubmodule.uniform_row_inputs(_sub(_Sampler(masters=False)), "decode") is None + # a row without loopback in the batch: row by row + assert LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=False) is None + assert LLMSubmodule.uniform_row_inputs(sub, "prefill_text", has_loopback=True) is None + assert LLMSubmodule.uniform_row_inputs( + _sub(_Sampler(masters=False)), "decode", has_loopback=True, + ) is None monkeypatch.setenv("MSTAR_DEVICE_LOOPBACK", "0") - assert LLMSubmodule.uniform_row_inputs(sub, "decode") is None + assert LLMSubmodule.uniform_row_inputs(sub, "decode", has_loopback=True) is None diff --git a/test/modular/test_sampler_last_token.py b/test/modular/test_sampler_last_token.py index b3a6b85d30..69845b47b4 100644 --- a/test/modular/test_sampler_last_token.py +++ b/test/modular/test_sampler_last_token.py @@ -165,3 +165,17 @@ def test_ingraph_scatter_masters_never_grow(): plain.register_request(rid) plain.write_last_tokens(["c"], torch.tensor([7])) # a gpu-thread access grows the device rows assert plain.last_token.master.shape[0] >= 3 and not plain.ingraph_scatter + + +def test_capture_padding_rows_take_no_slot(): + bufs = _buffers() + free = len(bufs._free_slots) + for rid in (-2, -3, -4): + bufs.register_request(rid) + assert len(bufs._free_slots) == free and not bufs.has_slot(-2) + bufs.register_request("a") + bufs.write_last_tokens(["a"], torch.tensor([5])) + # padding rows read slot 0 + assert bufs.last_tokens_for([-2, "a"]).tolist() == [ + bufs.last_token.master[0].item(), 5, + ] diff --git a/test/modular/test_worker_launch_signal.py b/test/modular/test_worker_launch_signal.py deleted file mode 100644 index 1c470e59f9..0000000000 --- a/test/modular/test_worker_launch_signal.py +++ /dev/null @@ -1,67 +0,0 @@ -"""The submitter's launch signal: set by the engine at the replay launch by -default, or by the gpu thread after commit and staging when the knob is on.""" - -import threading -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") - -from mstar.model.submodule_base import BatchedModelOutput # noqa: E402 -from mstar.worker.worker import Worker # noqa: E402 - - -class _Engine: - def __init__(self): - self.seen_event = "unset" - - def exec_and_postprocess(self, node_batch): - # what the engine would hand to the replay as its launch signal - self.seen_event = node_batch.launch_started_event - return BatchedModelOutput() - - -class _Span: - def __enter__(self): - return self - - def __exit__(self, *exc): - return False - - -def _worker(deferred): - w = object.__new__(Worker) - w._launch_signal_after_commit = deferred - w._span = lambda name: _Span() - return w - - -def test_by_default_the_engine_gets_the_event_to_set_at_the_launch(): - engine = _Engine() - event = threading.Event() - batch = SimpleNamespace(launch_started_event=event, completion_event=None) - Worker._exec_step(_worker(False), engine, batch, None) - assert engine.seen_event is event - assert not event.is_set(), "the stub engine never launched, so nobody set it" - - -def test_with_the_knob_the_gpu_thread_signals_after_the_step(): - engine = _Engine() - event = threading.Event() - batch = SimpleNamespace(launch_started_event=event, completion_event=None) - Worker._exec_step(_worker(True), engine, batch, None) - assert engine.seen_event is None, "the engine must not signal at the launch" - assert event.is_set() and batch.launch_started_event is event - - -def test_a_raising_step_still_releases_the_submitter(): - class _Raising(_Engine): - def exec_and_postprocess(self, node_batch): - raise RuntimeError("forward failed") - - event = threading.Event() - batch = SimpleNamespace(launch_started_event=event, completion_event=None) - with pytest.raises(RuntimeError): - Worker._exec_step(_worker(True), _Raising(), batch, None) - assert event.is_set() and batch.launch_started_event is event From da9522670cf7136b8b95e795746867ac536f9453 Mon Sep 17 00:00:00 2001 From: NSagan271 Date: Wed, 14 Oct 2026 06:45:48 +0000 Subject: [PATCH 218/219] ruff --- mstar/model/submodule_base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index efbe85bbc5..238b6037f5 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -726,7 +726,8 @@ def device_loopback_signals(self, graph_walk: str) -> frozenset[str]: ``preprocess`` and every fallback batch arrives with the signal empty. ``prepare_inputs`` is told per row (``has_loopback``) whether its last step here kept the value on the device; a row without it - (e.g. its previous step ran on another rank) brings the tensor. On a walk whose output signals all stay on the device the + (e.g. its previous step ran on another rank) brings the tensor. + On a walk whose output signals all stay on the device the engine builds no per-request output dicts and does not call ``postprocess``: the stop check, the inline emit and the next step's speculation read the step's rows instead. Empty by default; From 4a08df8097d673e97e02d1727b9eeb846b4bbb39 Mon Sep 17 00:00:00 2001 From: NSagan271 Date: Wed, 14 Oct 2026 06:58:37 +0000 Subject: [PATCH 219/219] fix rows_only for a request that runs _exec_single but still doesn't have forward_batched --- mstar/engine/engine.py | 5 +++-- test/modular/test_engine_rows_only_tail.py | 9 ++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 9e0d578119..fedddc475f 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -940,9 +940,10 @@ def exec( # admit/plan/commit assume the whole batch reaches the forward # in order; an unbatchable walk with >1 request can't, so each # request runs its own full cycle and merges. - if not batch.running_batched and len(batch.request_ids) > 1 and lease is None: - # per-request outputs: the step cannot stay on the rows + if not batch.running_batched and lease is None: + # the unbatched forward has per-request outputs: no rows to stay on batch.rows_only = False + if not batch.running_batched and len(batch.request_ids) > 1 and lease is None: batch.outputs = self._exec_per_request(batch) else: batch.outputs = self._exec_single(batch) diff --git a/test/modular/test_engine_rows_only_tail.py b/test/modular/test_engine_rows_only_tail.py index adfb7961bc..2663183f6a 100644 --- a/test/modular/test_engine_rows_only_tail.py +++ b/test/modular/test_engine_rows_only_tail.py @@ -1,5 +1,6 @@ """A rows-only step runs only a submodule's batch tail, and only when it has -one; a step that falls back to the per-request path is not rows-only.""" +one; a step on the unbatched forward (per request or a single row) is not +rows-only.""" import threading from types import SimpleNamespace @@ -55,7 +56,8 @@ def test_rows_only_runs_the_batch_tail_only_when_overridden(sub_cls): assert len(eng.per_request_tails) == 1 -def test_the_per_request_fallback_clears_rows_only(): +@pytest.mark.parametrize("rids", [[1, 2], [1]]) +def test_an_unbatched_step_clears_rows_only(rids): eng = _engine(_Plain.__new__(_Plain)) eng._enable_nvtx = False eng._enable_profile = False @@ -69,8 +71,9 @@ def _per_request(batch): return BatchedModelOutput() eng._exec_per_request = _per_request + eng._exec_single = _per_request batch = SimpleNamespace( - node_name="llm", request_ids=[1, 2], rows_only=True, running_batched=False, + node_name="llm", request_ids=rids, rows_only=True, running_batched=False, step_context=SimpleNamespace(slot_lease=None, graph_walk="decode"), exec_timings=SimpleNamespace(start=None), admit_error=None, outputs_ready=threading.Event(), release_waiters=lambda: None,