perf: sampler loop-back masters per node (fixes Whisper), KV plan cache, page-locked index arrays, in-graph decode knobs

docs/environment_variables.rst

@@ -329,15 +329,58 @@ 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_KV_CHAIN_LAZY_STEPS``
+     - ``1``
+     - How many decode steps of sampled ids the KV manager holds back before
+       it extends the requests' prefix chains (the per-page keys the prefix
+       cache is indexed by), one extend per request per batch of steps instead
+       of one per step. The per-request Python of that extension was 0.3 to
+       0.4 ms a step at 128 rows. A page key then appears up to that many
+       steps late and a later commit indexes it, as the single-request path
+       already allows; a request that finishes in between may leave its last
+       page unindexed. The default ``1`` extends on every step: ``16``
+       measured as noise at c64/c128 on 0.8B (the extension is a small part
+       of the stop check's tail).
+   * - ``MSTAR_KV_PINNED_INDPTRS``
+     - ``0``
+     - The KV plan writes FlashInfer's index arrays (qo_indptr, kv_indptr,
+       kv_indices, last_page_len) into a ring of page-locked host buffers and
+       the decode wrapper copies the indices straight into its captured
+       device buffer from there, so the copies FlashInfer's ``plan`` issues
+       are really asynchronous. From pageable memory the driver stages each
+       one and FlashInfer makes the indices copy blocking. Measured neutral
+       to slightly negative at c32-c128 on 0.8B (the event bookkeeping per
+       take costs about what it saves), so it is off by default.
+   * - ``MSTAR_KV_PLAN_CACHE``
+     - ``0``
+     - The KV manager plans a captured decode step off the previous step's
+       plan when the batch is the same rows in the same order and nothing
+       but that step's commit touched the streams: every length is one
+       more, pages change only where a row crossed a page boundary, and the
+       five FlashInfer index arrays come from numpy instead of a per-row
+       Python pass. The first version measured a loss at c64/c128 on 0.8B
+       (batches change every step or two there, and a hit still rebuilt the
+       views); the reworked path keeps the views lazy and is off by default
+       until measured. ``0`` plans every step from the streams. ``2`` runs
+       both and logs a difference, for checking a run.
    * - ``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.
+       cleanup per step). Only nodes that declare the read
+       (``reads_device_loopback``: Qwen3.5 decode) and only when the sampler
+       has graph buffers; the sampler keeps the masters for those nodes
+       alone, so a node that takes its tokens on the host (Whisper's
+       decoder) no longer pays the eager-step master write, a small
+       synchronous host-to-device copy per step. ``0`` keeps every node on
+       the per-request tensor path.
+   * - ``MSTAR_INGRAPH_DECODE_ROPE``
+     - ``0``
+     - A captured decode step builds its rotary tables inside the graph from
+       the position resource's planned positions instead of staging them.
+       Experiment knob, measured and gated on its own.
    * - ``MSTAR_SAMPLER_INGRAPH_SCATTER``
      - ``0``
      - The captured sampler gathers its RNG offsets from the slot masters

mstar/engine/cuda_graph_runner.py

@@ -223,7 +223,7 @@ def ensure(self, key: str, bs: int) -> list[int]:
         held = self._held.setdefault(key, [])
         rids = self.handles(key, bs)
         for rid in rids[len(held):]:
-            self._step_runner.ingest_request(rid)
+            self._step_runner.ingest_capture_row(rid)
             held.append(rid)
         return rids
 
@@ -1271,7 +1271,7 @@ def _capture_one(self, shape: PiecewiseCaptureShape, slot: int) -> None:
             dummy_rids, shape.seq_lens, self._bucket(shape), capture=True,
             slot=slot
         )
-        call = self._call_inputs(static_inputs, dummy_rids)
+        call = self._call_inputs(static_inputs, dummy_rids, step)
 
         fn = self._config.capture_fn
         if self._config.compile:
@@ -1319,19 +1319,24 @@ def _call_inputs(
         self,
         static_inputs: dict[str, torch.Tensor],
         step_ids: list[int],
+        step=None,
     ) -> PiecewiseCallInputs:
         """What the region is handed, over the padded capture batch.
 
         Built once, for the capture: the region's Python runs only there, so
         the rows are the runner's padding ids and the only thing a replay
-        changes is the contents of the static buffers.
+        changes is the contents of the static buffers. ``step`` is the
+        capture's declaration, so the region can tell it runs on a leased
+        slot (its context names the lease) and record the in-graph reads of
+        the slot's buffers a replay relies on.
         """
         return PiecewiseCallInputs(
             static_inputs=static_inputs,
             engine_inputs=ModelInputsFromEngine(
                 request_ids=list(step_ids),
                 per_request_info=dummy_metadata(step_ids, PIECEWISE_WALK),
                 resources=dict(self._resources),
+                step=step,
             ),
             kwargs=self._config.forward_kwargs,
         )

mstar/engine/engine.py

@@ -533,6 +533,13 @@ def load_model(
                 resources=resources
             )
             submodule.bind_node_resources(resources)
+            if getattr(submodule, "reads_device_loopback", False):
+                # the sampler keeps its slot masters only for a node that
+                # reads them back on the device (see NodeSubmodule)
+                for resource in resources.values():
+                    enable = getattr(resource, "enable_device_loopback_reader", None)
+                    if enable is not None:
+                        enable()
 
     def _compile_submodules(self) -> None:
         """Apply torch.compile to submodule forward paths.

mstar/engine/resources/attn/wrappers.py

@@ -298,6 +298,18 @@ def plan(
         Inputs may be on CPU; see prefill wrapper's plan docstring.
         """
         n_req = paged_kv_indptr.shape[0] - 1
+        if (
+            self.use_cuda_graph and paged_kv_indices.device.type != "cuda"
+            and paged_kv_indices.is_pinned()
+        ):
+            # FlashInfer copies host indices with a blocking copy; page-locked
+            # indices (the KV plan's ring, MSTAR_KV_PINNED_INDPTRS) go in one
+            # non-blocking copy into its own buffer instead, and its copy_ of
+            # that buffer onto itself is a no-op. Pageable indices keep
+            # FlashInfer's path, so the knob is a full toggle.
+            n = paged_kv_indices.shape[0]
+            self._paged_kv_indices_buf[:n].copy_(paged_kv_indices, non_blocking=True)
+            paged_kv_indices = self._paged_kv_indices_buf[:n]
 
         if self.enable_nvtx:
             from mstar.utils.profiler import range_pop, range_push

mstar/engine/resources/base.py

@@ -150,6 +150,10 @@ def fingerprint(self) -> bytes | None:
     def ingest_request(self, rid: str, overrides: ResourceReqConfig | None):
         return
 
+    def ingest_capture_row(self, rid: str):
+        """Open a CUDA graph capture padding row; by default, as a request."""
+        self.ingest_request(rid, None)
+
     def remove_request(self, rid: str):
         return