[GDN, Attention] Make a mixed decode + prompt step cheaper

docs/adding_models.rst

@@ -518,7 +518,7 @@ Step 2a — Declare your resources
 --------------------------------
 
 ``get_node_resources`` returns a flat list of ``NodeResourceSpec`` objects. Every spec has
-three common fields:
+these common members:
 
 - ``resource_key`` is the name of this resource. Three places use this name, and all three
   must agree: the layers bind to it, ``declare_step`` uses it as the key of its
@@ -535,6 +535,11 @@ three common fields:
   name. If a spec names a dependency that the model does not declare, loading fails
   immediately, rather than during a forward pass. ``RaggedAttentionSpec`` is the one
   attention spec that depends on nothing: it names no cache, because it has none.
+- ``combined_splits()`` returns a ``CombinedSplit`` for each combined walk this spec
+  runs on split kernels (see :ref:`combined walks <combined-walks>`); the default is
+  none. The engine reads only this hook to order a split walk's rows, so a spec whose
+  config carries ``splits`` must return them here. ``AttentionSpec`` and
+  ``LinearAttnSpec`` return their config's ``splits``.
 
 The spec types are:
 
@@ -1459,6 +1464,8 @@ 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:
+
 **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
@@ -1467,7 +1474,15 @@ can key on. A step of one walk keeps that walk's cap. A step that mixes walks ta
 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``).
+step may hold, for kernels sized to a few prompt rows (see ``CombinedSplit``). To run a
+combined walk on such split kernels, pass a ``CombinedSplit`` in the ``splits`` of each
+resource config that supports one (``AttentionConfig`` on the ``FLASHINFER`` backend, for
+captured steps only, and ``LinearAttnConfig`` with the GDN variant; other backends ignore
+``splits`` and run the walk unsplit); their specs return it from ``combined_splits()``,
+which is what the engine reads, and resources of one node must agree on a walk's
+``prefill_walks``. The engine then orders that walk's steps with the ``prefill_walks``
+rows last, so the kernels read the decode rows as the step's first tokens, in place. Every
+other walk of the combined walk must be one token per row.
 
 ``ARNodeSubmodule`` declares ``preprocess`` abstract because autoregressive nodes are
 expected to collate a batch. You can still disable batching for one node, or for specific

mstar/engine/engine.py

@@ -489,6 +489,8 @@ def __init__(
         self._resources: dict[str, Resource] = {}
         # resource key -> the nodes on this engine that use it
         self._resource_nodes: dict[str, frozenset[str]] = {}
+        # (node, combined walk) -> the walks whose rows a step of it runs last
+        self._split_prompt_walks: dict[tuple[str, str], frozenset[str]] = {}
         self._submodules: dict[str, SubmoduleManagement] = {}
         self._runner: StepRunner = None
         self._graph_runtime = graph_runtime
@@ -568,6 +570,7 @@ def load_model(
             for node in relevant_nodes:
                 node_to_resources.setdefault(node, []).append(spec.resource_key)
             self._resource_nodes[spec.resource_key] = frozenset(relevant_nodes)
+            self._add_split_row_order(spec, relevant_nodes)
             declared_by[spec.resource_key] = frozenset(spec.nodes)
 
         # rid -> (node, walk) of its last step here when that step was rows-only,
@@ -2284,6 +2287,33 @@ def capture_group(
             per_request_input_metadata={request_id: input_metadata},
         )
 
+    def _add_split_row_order(self, spec: NodeResourceSpec, nodes: set[str]) -> None:
+        """Record the row order ``spec``'s split kernels read (``CombinedSplit``)."""
+        for split in spec.combined_splits():
+            last = frozenset(split.prefill_walks)
+            for node in nodes:
+                key = (node, split.combined_walk)
+                if self._split_prompt_walks.setdefault(key, last) != last:
+                    raise ValueError(
+                        f"resource {spec.resource_key!r}: {node!r} splits combined walk "
+                        f"{split.combined_walk!r} with prefill walks {sorted(last)}, "
+                        f"another resource with {sorted(self._split_prompt_walks[key])}"
+                    )
+
+    def order_split_rows(
+        self, node_name: str, graph_walk: str, request_ids: list[int],
+        request_walks: Mapping[int, str],
+    ) -> list[int]:
+        """``request_ids`` as a step of ``graph_walk`` runs them: a walk run on
+        split kernels takes its prefill walks' rows last, in their order."""
+        last = self._split_prompt_walks.get((node_name, graph_walk))
+        if last is None:
+            return request_ids
+        first = [rid for rid in request_ids if request_walks.get(rid) not in last]
+        if len(first) == len(request_ids):
+            return request_ids
+        return first + [rid for rid in request_ids if request_walks.get(rid) in last]
+
     @staticmethod
     def _max_rows_in_combined(
         resource_key: str,

mstar/engine/resources/attn/config.py

@@ -36,9 +36,10 @@ class AttentionConfig:
     # pass the *original* qk head dim's scale, not the latent width's. Only the MLA
     # backend reads it; the others refuse it at build
     sm_scale: float | None = None
-    # FlashInfer only: combined walks whose captured steps plan their
-    # one-token rows on the decode wrapper and their prompt rows on a prefill
-    # wrapper sized to them (`FlashInferMixedWrapper`)
+    # FlashInfer only: combined walks whose captured steps run their leading
+    # one-token rows on a decode half and their prompt rows on a prefill
+    # wrapper sized to them (`FlashInferMixedWrapper`); the engine orders
+    # every step of those walks decode first
     splits: tuple[CombinedSplit, ...] = ()
     # FlashInfer only: run one-token decode rows on the tensor-core kernel.
     # False picks the CUDA-core decode kernel, which can win for MHA and
@@ -63,6 +64,9 @@ def depends_on(self) -> set[str]:
         # be the same KVConfig the KV resource holds
         return {self.config.kv_cache}
 
+    def combined_splits(self) -> tuple[CombinedSplit, ...]:
+        return self.config.splits
+
     @property
     def resource_class(self) -> "type[Resource]":
         from mstar.engine.resources.attn.base import AttentionManager

mstar/engine/resources/attn/wrappers.py

@@ -416,17 +416,21 @@ def run(self, q: torch.Tensor, kv_cache_layer: torch.Tensor) -> torch.Tensor:
 class FlashInferMixedWrapper:
     """A step mixing one-token rows (decode) with prompt rows, on two plans.
 
-    The packed prefill planner over every row costs per row about twice what
-    the decode planner does, and its kernel runs the one-token rows as
-    prefill queries. Here the decode wrapper takes every row of the bucket
-    (a one-token row attends its own pages; any other row reads one token of
-    the sink page, and its output is dropped), and the prefill wrapper takes
-    at most ``max_prefill_rows`` prompt rows, their tokens moved to the front
-    by a planned permutation. Both are captured at static shapes: the
-    bucket's rows and tokens, and ``max_prefill_rows`` sequences.
-
-    Every host array is staged in page-locked buffers this wrapper owns: a
-    mixed step plans inline on the gpu thread, behind the step in flight.
+    The engine orders a split walk's rows decode first (``CombinedSplit``),
+    so the step's leading run of one-token rows is its first tokens, row i
+    token i. The decode half plans the bucket's ``batch_size`` rows: a row of
+    the run attends its own pages, and any other reads the sink page and its
+    output is dropped (with tensor cores those rows get no query, except the
+    last, which takes the rest of the half's ``batch_size`` tokens). The
+    prefill wrapper takes every token: the run as one row on the sink page,
+    at most ``max_prefill_rows`` prompt rows, and a sink-page tail row over
+    the bucket's padding. Both read the step's queries in place, at static
+    shapes, and the decode outputs overwrite the run's.
+
+    A mixed step plans inline on the gpu thread, behind the step in flight,
+    so every host array is staged page-locked: this wrapper's own and the
+    CUDA-core decode half's in its buffers, the prefill wrapper's and the
+    tensor-core half's in theirs (``_stage_pinned``).
     """
 
     def __init__(
@@ -457,21 +461,33 @@ def __init__(
             num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads,
             head_dim=head_dim, page_size=page_size, device=device, backend=backend,
         )
-        self._decode = FlashInferDecodeWrapper(
-            workspace_buffer=decode_workspace, batch_size=batch_size,
-            max_num_pages=max_num_pages, accelerator_graph=True,
-            use_tensor_cores=use_tensor_cores, **kv,
-        )
-        # one row past the prompt rows takes the step's other tokens (see `plan`)
+        self._tensor_cores = use_tensor_cores
+        if use_tensor_cores:
+            # The kernel FlashInfer's tensor-core decode runs (FA2 paged
+            # prefill), planned with this step's qo_indptr: its KV split budget
+            # (2 CTAs per SM over the KV heads) charges each row with a query,
+            # so the rows outside the run take none, where the decode wrapper
+            # gives each one token. One row carries the outside rows' tokens.
+            self._decode = FlashInferPrefillWrapper(
+                workspace_buffer=decode_workspace, batch_size=batch_size,
+                max_total_tokens=batch_size, max_num_pages=max_num_pages,
+                accelerator_graph=True,
+                **{**kv, "backend": "fa2" if backend == "auto" else backend},
+            )
+        else:
+            self._decode = FlashInferDecodeWrapper(
+                workspace_buffer=decode_workspace, batch_size=batch_size,
+                max_num_pages=max_num_pages, accelerator_graph=True,
+                use_tensor_cores=False, **kv,
+            )
+        # a row ahead of the prompt rows takes the decode run's tokens, and one
+        # behind them the bucket's padding (see `plan`)
         self._prefill = FlashInferPrefillWrapper(
-            workspace_buffer=prefill_workspace, batch_size=max_prefill_rows + 1,
+            workspace_buffer=prefill_workspace, batch_size=max_prefill_rows + 2,
             max_total_tokens=max_total_tokens, max_num_pages=max_num_pages,
             accelerator_graph=True, **kv,
         )
-        i64 = dict(dtype=torch.int64, device=device)
-        self._perm = torch.zeros(max_total_tokens, **i64)
-        self._dec_src = torch.zeros(batch_size, **i64)
-        self._dec_dst = torch.zeros(batch_size, **i64)
+        self._dec_mask = torch.zeros(batch_size, dtype=torch.bool, device=device)
         # the whole step's qo_indptr, which `select_last_hidden` reads
         self._qo_indptr_buf = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
         self._pinned: dict[str, torch.Tensor] = {}
@@ -525,16 +541,21 @@ def plan(
         spans = np.diff(qo)
         rows = len(spans)
         tokens = max(self._tokens, int(qo[-1]))
-        dec = spans == 1  # a one-token row is a decode query, state or not
-        prompt = np.flatnonzero(spans > 1)
+        one = spans == 1  # a one-token row is a decode query, state or not
+        run = rows if one.all() else int(np.argmin(one))
+        prompt = run + np.flatnonzero(spans[run:] > 0)
         if len(prompt) > self._max_prefill_rows:
             raise RuntimeError(
                 f"FlashInferMixedWrapper: {len(prompt)} prompt rows in a step "
-                f"captured for {self._max_prefill_rows}; the scheduler caps them"
+                f"captured for {self._max_prefill_rows}. Possible causes: the "
+                "scheduler let in more than max_rows_in_combined allows, a decode "
+                "row behind a prompt row (a broken row order), or a CombinedSplit "
+                "that leaves out a multi-token walk"
             )
+        dec = np.arange(rows) < run
 
-        # decode: every row of the bucket; a row that is not one token reads
-        # one token of the sink page, and its output lands on a scratch row
+        # decode: every row of the bucket; a row outside the run reads one
+        # token of the sink page, and its output is dropped
         cnt = np.where(dec, npages, 1)
         d_indptr = np.zeros(rows + 1, dtype=np.int32)
         np.cumsum(cnt, out=d_indptr[1:])
@@ -546,54 +567,62 @@ def plan(
         d_last = np.where(dec, last, 1).astype(np.int32)
         d_lens = np.where(dec, lens, 1).astype(np.int32)
 
-        # prefill: the prompt rows, then a sink-page tail row over the step's
-        # other tokens, so qo_indptr ends at the token count FlashInfer checks
+        # prefill: a sink-page row over the run's tokens, the prompt rows,
+        # empty padding rows, then a sink-page row over the bucket's padding,
+        # so qo_indptr ends at the token count FlashInfer checks
         pad = self._max_prefill_rows - len(prompt)
-        prompt_tokens = int(spans[prompt].sum())
-        tail = tokens - prompt_tokens
-        tail_kv = max(tail, 1)
-        tail_pages = -(-tail_kv // self._page_size)
-        p_spans = np.concatenate([spans[prompt], np.zeros(pad, dtype=np.int64), [tail]])
-        p_qo = np.zeros(self._max_prefill_rows + 2, dtype=np.int32)
+        tail = tokens - run - int(spans[prompt].sum())
+
+        def sink_row(n):
+            n = max(n, 1)
+            pages = -(-n // self._page_size)
+            return pages, n - (pages - 1) * self._page_size, n
+
+        lead_pages, lead_last, lead_kv = sink_row(run)
+        tail_pages, tail_last, tail_kv = sink_row(tail)
+        p_spans = np.concatenate([[run], spans[prompt], np.zeros(pad, dtype=np.int64), [tail]])
+        p_qo = np.zeros(self._max_prefill_rows + 3, dtype=np.int32)
         np.cumsum(p_spans, out=p_qo[1:])
-        p_cnt = np.concatenate([npages[prompt], np.ones(pad, dtype=np.int64), [tail_pages]])
-        p_kvp = np.zeros(self._max_prefill_rows + 2, dtype=np.int32)
+        p_cnt = np.concatenate([[lead_pages], npages[prompt], np.ones(pad, dtype=np.int64), [tail_pages]])
+        p_kvp = np.zeros(self._max_prefill_rows + 3, dtype=np.int32)
         np.cumsum(p_cnt, out=p_kvp[1:])
         p_idx = np.concatenate(
-            [idx[kvp[r]:kvp[r + 1]] for r in prompt]
+            [np.full(lead_pages, self._sink_page, dtype=idx.dtype)]
+            + [idx[kvp[r]:kvp[r + 1]] for r in prompt]
             + [np.full(pad + tail_pages, self._sink_page, dtype=idx.dtype)]
         ).astype(np.int32)
-        tail_last = tail_kv - (tail_pages - 1) * self._page_size
-        p_last = np.concatenate([last[prompt], np.ones(pad), [tail_last]]).astype(np.int32)
-        p_lens = np.concatenate([lens[prompt], np.ones(pad), [tail_kv]]).astype(np.int32)
+        p_last = np.concatenate([[lead_last], last[prompt], np.ones(pad), [tail_last]]).astype(np.int32)
+        p_lens = np.concatenate([[lead_kv], lens[prompt], np.ones(pad), [tail_kv]]).astype(np.int32)
 
-        # the prompt tokens first, then every other position in order
-        front = np.concatenate(
-            [np.arange(qo[r], qo[r + 1]) for r in prompt] or [np.zeros(0, dtype=np.int64)]
-        )
-        rest = np.ones(tokens, dtype=bool)
-        rest[front] = False
-        perm = np.concatenate([front, np.flatnonzero(rest)])
-        dec_src = np.minimum(qo[:-1], max(tokens - 1, 0))
-        dec_dst = np.where(dec, qo[:-1], tokens)
-
-        for name, values, dst in (
-            ("perm", perm, self._perm), ("dec_src", dec_src, self._dec_src),
-            ("dec_dst", dec_dst, self._dec_dst),
-        ):
-            dst[: len(values)].copy_(self._pin(name, values, torch.int64), non_blocking=True)
+        self._dec_mask[:rows].copy_(self._pin("dec_mask", dec, torch.bool), non_blocking=True)
         self._qo_indptr_buf[: rows + 1].copy_(
             self._pin("qo", qo.astype(np.int32), torch.int32), non_blocking=True,
         )
         self._num_rows = rows
 
-        self._decode.plan(
-            paged_kv_indptr=self._pin("d_indptr", d_indptr, torch.int32),
-            paged_kv_indices=self._pin("d_indices", d_indices, torch.int32),
-            paged_kv_last_page_len=self._pin("d_last", d_last, torch.int32),
-            kv_lens=self._pin("d_lens", d_lens, torch.int32).as_subclass(HostLens),
-            dtype=dtype, window_left=window_left,
-        )
+        if self._tensor_cores:
+            # the run's rows one token each, and the last row the rest, so
+            # qo_indptr ends at the row count FlashInfer checks
+            d_qo = np.zeros(rows + 1, dtype=np.int32)
+            d_qo[1 : run + 1] = np.arange(1, run + 1)
+            d_qo[run + 1 :] = run
+            d_qo[-1] = rows
+            self._decode.plan(
+                qo_indptr=torch.from_numpy(d_qo),
+                paged_kv_indptr=torch.from_numpy(d_indptr),
+                paged_kv_indices=torch.from_numpy(d_indices),
+                paged_kv_last_page_len=torch.from_numpy(d_last),
+                kv_lens=torch.from_numpy(d_lens).as_subclass(HostLens),
+                causal=False, dtype=dtype, window_left=window_left,
+            )
+        else:
+            self._decode.plan(
+                paged_kv_indptr=self._pin("d_indptr", d_indptr, torch.int32),
+                paged_kv_indices=self._pin("d_indices", d_indices, torch.int32),
+                paged_kv_last_page_len=self._pin("d_last", d_last, torch.int32),
+                kv_lens=self._pin("d_lens", d_lens, torch.int32).as_subclass(HostLens),
+                dtype=dtype, window_left=window_left,
+            )
         self._prefill.plan(
             qo_indptr=torch.from_numpy(p_qo),
             paged_kv_indptr=torch.from_numpy(p_kvp),
@@ -606,12 +635,12 @@ def plan(
 
     @torch.compiler.disable
     def run(self, q: torch.Tensor, kv_cache_layer: torch.Tensor) -> torch.Tensor:
-        n = q.shape[0]
-        front = q.index_select(0, self._perm[:n])
-        rows = q.index_select(0, self._dec_src[: self._num_rows])
-        out_front = self._prefill.run(front, kv_cache_layer)
+        n = self._num_rows
+        rows = q[:n] if q.shape[0] >= n else torch.cat((q, q.new_zeros((n - q.shape[0], *q.shape[1:]))))
+        out = self._prefill.run(q, kv_cache_layer)
         out_rows = self._decode.run(rows, kv_cache_layer)
-        out = out_front.new_empty((n + 1, *out_front.shape[1:]))
-        out.index_copy_(0, self._perm[:n], out_front[:n])
-        out.index_copy_(0, self._dec_dst[: self._num_rows], out_rows.to(out_front.dtype))
-        return out[:n]
+        m = min(n, out.shape[0])
+        out[:m] = torch.where(
+            self._dec_mask[:m].view(-1, 1, 1), out_rows[:m].to(out.dtype), out[:m],
+        )
+        return out[: q.shape[0]]

mstar/engine/resources/linear_attn/config.py

@@ -64,6 +64,9 @@ class LinearAttnSpec(NodeResourceSpec):
     def depends_on(self) -> set[str]:
         return {self.config.recurrent_state}
 
+    def combined_splits(self) -> tuple[CombinedSplit, ...]:
+        return self.config.splits
+
     @property
     def writes_prefix_snapshots(self) -> bool:
         # only the GDN wrapper writes state checkpoints into the pool's snapshot slots

mstar/engine/resources/linear_attn/gdn.py

@@ -202,7 +202,10 @@ def _get_wrapper(
         else:
             key = CGSlotKey(lease.bucket, lease.slot, label)
             store = self._cg_wrappers
-            bs = max(self._cg_max_bs, lease.bucket.bs)
+            # A split's decode half runs on as many rows as it is sized for,
+            # so it takes the bucket's; the floor is for the other wrappers,
+            # whose kernels run on the step's rows whatever their buffers hold.
+            bs = lease.bucket.bs if split is not None else max(self._cg_max_bs, lease.bucket.bs)
             tok = lease.bucket.num_tokens
 
         if key not in store: