[Chunked Prefill #2] Combined Walks

docs/adding_models.rst

@@ -966,6 +966,16 @@ into one GPU call. A submodule controls this behavior with three methods:
   ``forward``.
 - ``max_batch_size(self, graph_walk)`` sets an optional upper limit.
 
+**Combined walks.** By default a step holds one walk. A model lets several walks of a node
+share a step by returning ``{node: {combined_walk: {walks}}}`` from
+``get_combined_graph_walks``, for example decode rows beside a prompt's prefill. The
+combined walk is a name of its own, which ``max_batch_size`` and the CUDA-graph configs
+can key on. A step of one walk keeps that walk's cap. A step that mixes walks takes the
+combined walk's ``max_batch_size`` when it has one, else the smallest cap of the walks it
+holds, so a small prefill cap never limits a decode-only step.
+``max_rows_in_combined(self, graph_walk)`` also bounds how many rows of one walk a mixed
+step may hold, for kernels sized to a few prompt rows (see ``CombinedSplit``).
+
 ``ARNodeSubmodule`` declares ``preprocess`` abstract because autoregressive nodes are
 expected to collate a batch. You can still disable batching for one node, or for specific
 walks, by returning ``False`` from ``can_batch`` in those cases.

docs/serving.rst

@@ -184,6 +184,14 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
        conductor's queue.
    * - ``model_kwargs``
      - *(optional)* Server-init model parameters (see below).
+   * - ``speculation_mixes_walks``
+     - *(optional)* Per node, ``true``/``false`` for every walk of the node's combined
+       walks, or a mapping of walk -> ``true``/``false``. On for a walk, a speculation
+       chain on it merges ready rows of the other walks of its combined walk, so
+       decode steps take in new prompts as they arrive. Off (the default), a chain
+       keeps to its own walk and yields to the others, and a step mixes walks only
+       when the scheduler builds one fresh. Naming a node with no combined walk, or a
+       walk outside its combined walks, fails loading.
    * - ``warmup_requests``
      - *(optional)* Requests the server runs after the workers are ready and before it
        binds, so no client pays a shape's first-request cost (torch.compile of a step,

mstar/engine/cuda_graph_config.py

@@ -125,6 +125,15 @@ def distribute_tokens(total_tokens: int, bs: int) -> list[int]:
     return seq_lens
 
 
+def decode_rows_then_prompt(total_tokens: int, bs: int) -> list[int]:
+    """``bs - 1`` one-token rows and a prompt row of the rest: a mixed step's
+    layout, for kernels that split rows by span. `distribute_tokens` when
+    there are fewer tokens than rows."""
+    if total_tokens < bs:
+        return distribute_tokens(total_tokens, bs)
+    return [1] * (bs - 1) + [total_tokens - (bs - 1)]
+
+
 class PackedCudaGraphConfig(CudaGraphConfig):
     def __init__(
         self,
@@ -144,6 +153,10 @@ def __init__(
         # the resources size their per-bucket buffers for the whole plan.
         total_tokens_multiplier: int = 1,
         input_seq_dims: dict[str, int] | None = None,
+        # (input tokens, bs) -> each capture row's span; `distribute_tokens`
+        # by default. A walk whose kernels split rows by span (a mixed step's
+        # one-token rows beside a prompt) captures a layout it will replay.
+        token_layout: Callable[[int, int], list[int]] | None = None,
     ):
         super().__init__(
             capture_graph_walk=capture_graph_walk,
@@ -160,16 +173,22 @@ def __init__(
         self.make_node_input = make_node_input
         self.capture_token_lengths = capture_token_lengths
         self.total_tokens_multiplier = total_tokens_multiplier
+        self.token_layout = token_layout or distribute_tokens
 
     def get_config_type(self) -> CudaGraphConfigType:
         return CudaGraphConfigType.FLASH_INFER_PACKED
 
     def get_total_tokens(self, bs: int) -> list[int]:
-        return [n * self.total_tokens_multiplier for n in self.capture_token_lengths]
+        # every real row carries a token, so a bucket with fewer tokens than
+        # rows could never serve a step; capturing it would give rows of none
+        return [
+            n * self.total_tokens_multiplier for n in self.capture_token_lengths
+            if n >= bs
+        ]
 
     def get_node_inputs(self, bs: int, num_tokens: int):
         # ``num_tokens`` is the bucket's plan tokens; the rows get input tokens
-        seq_lens = distribute_tokens(num_tokens // self.total_tokens_multiplier, bs)
+        seq_lens = self.token_layout(num_tokens // self.total_tokens_multiplier, bs)
         return [
             self.make_node_input(n) for n in seq_lens
         ]

mstar/engine/cuda_graph_runner.py

@@ -191,6 +191,10 @@ def _stop_recording_to_pool(device, pool) -> None:
         pass
 
 
+def _is_packed(config: CudaGraphConfig) -> bool:
+    return config.get_config_type() == CudaGraphConfigType.FLASH_INFER_PACKED
+
+
 class DummyRowPool:
     """The padding rows captured replays pad onto.
 
@@ -562,6 +566,7 @@ def _capture_one(self, spec: CGSlotSpec) -> CudaGraphSlot:
                 slot=spec.slot,
                 capture=True,
                 slot_lease=lease,
+                force_prefill=_is_packed(config),
             ))
 
         engine_inputs.step = step
@@ -826,6 +831,10 @@ def can_run(
             graph_walk, bs, num_tokens, cg_key_info
         ) is not None
 
+    def is_packed(self, lease: SlotLease) -> bool:
+        """Whether ``lease`` replays a packed capture."""
+        return _is_packed(self._buckets[lease.bucket].config)
+
     def lease_slot(
         self, graph_walk: str, bs: int,
         slot: int,
@@ -1366,6 +1375,7 @@ def _declare(
             slot=slot,
             capture=capture,
             slot_lease=SlotLease(slot=slot, bucket=bucket),
+            force_prefill=self._config.get_config_type() == PiecewiseConfigType.PACKED,
         )
         ctx.set_padded_rids(tuple(request_ids))
         step.set_ctx(ctx)

mstar/engine/engine.py

@@ -301,6 +301,11 @@ class ExecutingBatch:
     # the inline emit and the speculation splice read the rows instead.
     rows_only: bool = False
 
+    # A mixed step's rows whose own walk keeps every output signal on the
+    # device: such a row loops back as it would in a step of its walk alone,
+    # while the step's other rows route tensors.
+    loopback_rids: frozenset = frozenset()
+
     # 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)
@@ -497,6 +502,9 @@ def load_model(
                     dependencies={
                         key: specs_by_key[key] for key in spec.depends_on()
                     },
+                    max_rows_in_combined=self._max_rows_in_combined(
+                        spec.resource_key, relevant_nodes, submodules, model,
+                    ),
                 ),
             )
 
@@ -759,19 +767,24 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None:
         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);
-        # 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, ())
+        # read leniently, as the probe itself was only reached per prepared row.
+        # Under a combined walk each row has its own walk, so a mixed step
+        # probes, prepares and stamps per row, and takes no uniform path.
+        walk = batch.step_context.graph_walk
+        mixed = bool(batch.step_context.request_walks)
+        walk_of = batch.step_context.walk_of
+        keyed_walks = (getattr(self, "_keyed_walks", None) or {}).get(
+            batch.node_name, (),
+        )
         loopback = self._loopback_rows(batch)
-        if not probe_prefix:
+        if not mixed and walk not in keyed_walks:
             uniform = submodule.uniform_row_inputs(
                 walk, has_loopback=loopback is not None and all(loopback),
             )
@@ -785,10 +798,14 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None:
                     batch=batch, model_inputs=node_inputs
                 )
                 return
+        # a step of one walk probes (or not) for every row alike
+        step_probes = not mixed and walk in keyed_walks
         for i, rid in enumerate(batch.request_ids):
+            row_walk = walk_of(rid) if mixed else walk
+            probe_prefix = (row_walk in keyed_walks) if mixed else step_probes
             try:
                 req_inputs = prepare(
-                    graph_walk=walk,
+                    graph_walk=row_walk,
                     fwd_info=wrapped[rid],
                     inputs=tensors.get(rid, {}),
                     resources=resources,
@@ -799,12 +816,16 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None:
                     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)
+                if req_inputs is not None:
+                    if probe_prefix:
+                        req_inputs = self._skip_cached_prefix(batch, rid, req_inputs)
+                    if probe_prefix or mixed:
+                        # the submodule reads a mixed step's rows by their own walk
+                        req_inputs.graph_walk = row_walk
             except Exception as error:
                 logger.exception(
                     "prepare_inputs failed for request %s (node=%s, walk=%s)",
-                    rid, batch.node_name, batch.step_context.graph_walk,
+                    rid, batch.node_name, row_walk,
                 )
                 batch.register_failure(rid, error)
                 continue
@@ -825,8 +846,14 @@ def _loopback_rows(self, batch: ExecutingBatch) -> list[bool] | None:
         steps = self._loopback_steps
         if not steps:
             return None
-        key = (batch.node_name, batch.step_context.graph_walk)
         get = steps.get
+        ctx = batch.step_context
+        if ctx.request_walks:
+            # a mixed step: a row's token is on the device if its own walk's
+            # last step kept it there, whatever this step is labelled
+            node, walk_of = batch.node_name, ctx.walk_of
+            return [get(rid) == (node, walk_of(rid)) for rid in batch.request_ids]
+        key = (batch.node_name, ctx.graph_walk)
         return [get(rid) == key for rid in batch.request_ids]
 
     def _record_loopback(self, batch: ExecutingBatch) -> None:
@@ -835,6 +862,17 @@ def _record_loopback(self, batch: ExecutingBatch) -> None:
             steps.update(dict.fromkeys(
                 batch.request_ids, (batch.node_name, batch.step_context.graph_walk),
             ))
+        elif batch.loopback_rids:
+            # a mixed step: a row looped back under its own walk, so the next
+            # step of that walk finds its token on the device
+            keep, node, walk_of = (
+                batch.loopback_rids, batch.node_name, batch.step_context.walk_of,
+            )
+            for rid in batch.request_ids:
+                if rid in keep:
+                    steps[rid] = (node, walk_of(rid))
+                else:
+                    steps.pop(rid, None)
         elif steps:
             for rid in batch.request_ids:
                 steps.pop(rid, None)
@@ -866,18 +904,23 @@ def extend_prefix_chains(
         keys = self._runner.prefix_chain_keys(batch.node_name)
         if not keys:
             return
-        walk = batch.step_context.graph_walk
+        # The batch call keys each row under its own walk, so a combined-walk
+        # step takes it too; only the absence of host rows falls through.
         if host_rows is not None:
+            ctx = batch.step_context
             self._runner.extend_prefix_chains_batch(
-                batch.request_ids, batch.node_name, walk, host_rows, outputs,
-                keys=keys,
+                batch.request_ids, batch.node_name,
+                # a step of one walk passes the walk: no call per row
+                ctx.walk_of if ctx.request_walks else ctx.graph_walk,
+                host_rows, outputs, keys=keys,
             )
             return
         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, keys=keys,
+                    rid, batch.node_name, batch.step_context.walk_of(rid), per_rid,
+                    keys=keys,
                 )
 
     def _skip_cached_prefix(
@@ -888,7 +931,7 @@ def _skip_cached_prefix(
         Only the keyed walk is probed, and only when `split_inputs` can cut its
         inputs; a guided walk writes two labels from one input, so it is skipped.
         """
-        walk = batch.step_context.graph_walk
+        walk = batch.step_context.walk_of(rid)
         if walk not in self._keyed_walks.get(batch.node_name, ()):
             return inputs
         assert isinstance(inputs, ARNodeInputs) and inputs.custom_pos_ids is None, (
@@ -950,12 +993,14 @@ def exec(
                 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
+                    batch.loopback_rids = frozenset()
                 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)
             if batch.admit_error is None:
                 self._record_loopback(batch)
+                batch.outputs.loopback_rids = batch.loopback_rids
             batch.outputs_ready.set()
             return batch.outputs
         except Exception:
@@ -1483,8 +1528,10 @@ def check_stop_for_batch(
         submodule = self._submodules[batch.node_name].submodule
         if host_rows is not None:
             try:
+                ctx = batch.step_context
                 batched = submodule.check_stop_batched(
                     batch.request_ids, batch.per_request_info_wrapped, host_rows,
+                    step_walk=None if ctx.request_walks else ctx.graph_walk,
                 )
             except Exception:
                 logger.debug(
@@ -1533,13 +1580,15 @@ def finalize_batch(
                     rids,
                     node_name=batch.node_name,
                     graph_walk=batch.step_context.graph_walk,
+                    request_walks=batch.step_context.request_walks,
                 )
                 return batch.resource_publish_info
             # Returns rid -> {resource label -> published info}
             published = self._runner.publish(
                 rids,
                 node_name=batch.node_name,
                 graph_walk=batch.step_context.graph_walk,
+                request_walks=batch.step_context.request_walks,
             )
             for rid, info in batch.per_request_info.items():
                 if rid not in published:
@@ -1568,6 +1617,7 @@ def materialize_publish(
             snaps, todo,
             node_name=batch.node_name,
             graph_walk=batch.step_context.graph_walk,
+            request_walks=batch.step_context.request_walks,
         )
         for rid in todo:
             per_key = published.get(rid) or {}
@@ -1586,6 +1636,7 @@ def finalize_stopped_requests(
             request_ids,
             node_name=batch.node_name,
             graph_walk=batch.step_context.graph_walk,
+            request_walks=batch.step_context.request_walks,
         )
         for rid in request_ids:
             rid_published = published.get(rid, {})
@@ -1795,6 +1846,34 @@ def capture_group(
             graph_walk, {request_id: request_info},
         )
 
+    @staticmethod
+    def _max_rows_in_combined(
+        resource_key: str,
+        nodes: set[str],
+        submodules: dict[str, NodeSubmodule],
+        model: "Model | None",
+    ) -> dict[str, int | None]:
+        """Each walk of ``nodes`` to its ``max_rows_in_combined``, one value a
+        walk across the nodes a resource serves."""
+        if model is None:
+            return {}
+        caps: dict[str, int | None] = {}
+        for walk, graph in model.get_graph_walk_graphs().items():
+            for node in sorted(nodes & set(graph.get_nodes())):
+                cap = submodules[node].max_rows_in_combined(walk)
+                if walk in caps and caps[walk] != cap:
+                    raise ValueError(
+                        f"resource {resource_key!r}: the nodes it serves cap "
+                        f"walk {walk!r} differently in a combined step "
+                        f"({caps[walk]} and {cap})"
+                    )
+                caps[walk] = cap
+        return caps
+
+    def get_max_rows_in_combined(self, node_name: str, graph_walk: str) -> int | None:
+        """Rows of ``graph_walk`` a step mixing it with another walk may hold."""
+        return self._submodules[node_name].submodule.max_rows_in_combined(graph_walk)
+
     def get_max_batch_size(self, node_name: str, graph_walk: str) -> int | None:
         """Most requests this node will take in one step, or None for no cap.
 
@@ -1881,6 +1960,7 @@ def reserve_replay_slot(self, batch: ExecutingBatch) -> SlotLease | None:
         )
         if lease is not None:
             batch.lease_slot(lease)
+            batch.step_context.force_prefill = cg_runner.is_packed(lease)
         return lease
 
     def can_pre_plan(self, node_name: str) -> bool:

mstar/engine/resources/attn/base.py

@@ -96,6 +96,8 @@ def build(cls, spec: AttentionSpec, info: EngineResourceInfo):
                 dtype=info.kv_dtype,
                 kv_config=kv_config,
                 backend=spec.config.flashinfer_backend,
+                splits=spec.config.splits,
+                max_rows_in_combined=info.max_rows_in_combined,
             )
 
         if backend == AttnBackend.FLEX: