[KV cache] Fix multi-label plans under the host index ring and plan cache

docs/environment_variables.rst

@@ -348,9 +348,9 @@ Worker scheduling
        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.
+       one and FlashInfer makes the indices copy blocking. No measurable gain
+       at c32-c128 on 0.8B, so it is off by default. The cost is one CUDA
+       event per plan, which a buffer set waits on before it is rewritten.
    * - ``MSTAR_KV_PLAN_CACHE``
      - ``0``
      - The KV manager plans a captured decode step off the previous step's

mstar/engine/resources/kv/manager.py

@@ -348,10 +348,11 @@ class _DecodePlanCache:
     """The last leased decode plan, as `_fast_decode_plan` continues it."""
 
     __slots__ = (
-        "step_segments", "real_segments", "n_real", "pad_layout", "views", "epoch", "pending", "derived", "_rows",
+        "step_segments", "real_segments", "n_real", "pad_layout", "combined_labels", "views", "epoch", "pending",
+        "derived", "_rows",
     )
 
-    def __init__(self, step_segments, real_segments, n_real, pad_layout, views, epoch):
+    def __init__(self, step_segments, real_segments, n_real, pad_layout, combined_labels, views, epoch):
         # the planned step's own segments tuple: a commit of that object is the
         # one commit the next plan may assume
         self.step_segments = step_segments
@@ -360,6 +361,8 @@ def __init__(self, step_segments, real_segments, n_real, pad_layout, views, epoc
         self.n_real = n_real
         # (count, span) of the padding rows, which change dummy ids per slot
         self.pad_layout = pad_layout
+        # the plan labels the views were built for
+        self.combined_labels = combined_labels
         # the full plan's views per label, the source of `derived`
         self.views = views
         self.epoch = epoch
@@ -1542,7 +1545,7 @@ def _remember_decode_plan(
                 return
         self._decode_plan_cache = _DecodePlanCache(
             step_segments=step.segments, real_segments=step.segments[:n_real],
-            n_real=n_real, pad_layout=pad_layout,
+            n_real=n_real, pad_layout=pad_layout, combined_labels=dict(step.combined_labels),
             views={label: out.views for label, out in res.items()},
             epoch=self._plan_epoch,
         )
@@ -1561,7 +1564,12 @@ def _derive_rows(self, cache: _DecodePlanCache, ctx: StepContext) -> dict | None
                     break
                 n_real += 1
             real, pad = views[:n_real], views[n_real:]
-            if pad and (len(pad), pad[0].to_compute) != cache.pad_layout:
+            # a hit rebuilds each label's padding from all the step's padding
+            # rows; in a combined label, a later source label's real rows can
+            # follow an earlier one's padding
+            if [(v.request_id, v.label, v.to_compute) for v in pad] != [
+                (seg.request_id, seg.label, seg.span) for seg in cache.step_segments[cache.n_real:]
+            ]:
                 return None
             streams = []
             for v in real:
@@ -1614,7 +1622,13 @@ def _fast_decode_plan(
         if (
             len(segs) != n_real + cache.pad_layout[0]
             or segs[:n_real] != cache.real_segments
+            or step.combined_labels != cache.combined_labels
             or any(seg.span != cache.pad_layout[1] or ctx.is_real_row(seg.request_id) for seg in segs[n_real:])
+            # padding keeps its labels: their order places the real rows in a combined label
+            or any(
+                seg.label != prev.label
+                for seg, prev in zip(segs[n_real:], cache.step_segments[n_real:], strict=True)
+            )
         ):
             stats["miss"] += 1
             return None
@@ -1669,7 +1683,8 @@ def _fast_decode_plan(
             else:
                 last_all, npages_all, indices, pages_all = last, npages, rows.indices, rows.last_page.copy()
             kv_indptr = np.concatenate((np.zeros(1, np.int64), np.cumsum(npages_all))).astype(np.int32)
-            n_rows = n_real + n_pad
+            # this label's rows: `n_real` counts the real rows of every label
+            n_rows = len(rows.rows) + n_pad
             qo_indptr = np.arange(n_rows + 1, dtype=np.int32)
             kv_lens = (npages_all - 1) * page_size + last_all
             row_pages_now, lens_now = rows.row_pages, lens.tolist()
@@ -1803,6 +1818,8 @@ def plan(self, step: KVStep, ctx: StepContext) -> dict[str, KVPlanOutput]:
             pending = {}
             self._maybe_apply_forks(step, ctx)
 
+        if self._indptr_ring is not None:
+            self._indptr_ring.begin_plan()
         res = None
         if self._plan_cache_mode and not pending:
             res = self._fast_decode_plan(step, ctx)

mstar/engine/resources/kv/plan.py

@@ -142,60 +142,58 @@ class PinnedIndexRing:
     hands FlashInfer, so its ``copy_(non_blocking=True)`` calls are really
     asynchronous (from pageable memory the driver stages each one).
 
-    ``take`` fills the next buffer and returns a view of it. A buffer is
-    rewritten only after the copies out of it have completed: the stream a
-    plan ran on records an event when the next ``take`` happens on that
-    stream (by then the plan's copies were issued), and a slot waits on its
-    event before reuse. Plans do not overlap, so one pending slot suffices.
-    Without CUDA the buffers are plain tensors and nothing is recorded.
+    The ring holds ``depth`` buffer sets, one per plan. ``begin_plan`` moves to
+    the next set, and each ``take`` until the next ``begin_plan`` fills a buffer
+    of that set no other take uses (the set grows if needed), so every label
+    of a plan keeps its arrays (four takes each) however many labels there
+    are. A set is rewritten only after the copies out of it have completed:
+    the next ``begin_plan`` records an event on the stream the plan ran on (by
+    then the plan's copies, and its consumers', were issued), and a set waits
+    on its event before reuse. Plans do not overlap, so one pending set
+    suffices. Without CUDA the buffers are plain tensors and nothing is
+    recorded.
     """
 
-    def __init__(self, numel: int = 4096, depth: int = 8):
+    def __init__(self, numel: int = 4096, depth: int = 2, per_plan: int = 4):
         self._depth = depth
-        self._numel = 0
+        self._numel = max(numel, 256)
         self._pinned = torch.cuda.is_available()
-        self._bufs: list[torch.Tensor] = []
+        self._sets: list[list[torch.Tensor]] = [
+            [self._alloc(self._numel) for _ in range(per_plan)] for _ in range(depth)
+        ]
         self._events: list = [None] * depth
         self._streams: list = [None] * depth
-        self._next = 0
-        self._last: int | None = None
-        self._grow(numel)
-
-    def _grow(self, numel: int) -> None:
-        numel = max(numel, self._numel * 2, 256)
-        self._bufs = [
-            torch.empty(numel, dtype=torch.int32, pin_memory=self._pinned)
-            for _ in range(self._depth)
-        ]
-        self._numel = numel
+        self._set = 0
+        self._used = 0
 
-    def take(self, values) -> torch.Tensor:
-        n = len(values)
-        if n > self._numel:
-            self._grow(n)
+    def _alloc(self, numel: int) -> torch.Tensor:
+        return torch.empty(numel, dtype=torch.int32, pin_memory=self._pinned)
+
+    def begin_plan(self) -> None:
+        cur = self._set
+        nxt = (cur + 1) % self._depth
         if self._pinned:
-            stream = torch.cuda.current_stream()
-            if self._last is not None and self._streams[self._last] is not stream:
-                # the previous plan's copies were issued on another stream
-                # (an inline plan on the gpu thread after a pre-plan): record
-                # there, where they are
-                ev = self._events[self._last] or torch.cuda.Event()
-                self._events[self._last] = ev
-                ev.record(self._streams[self._last])
-            elif self._last is not None:
-                ev = self._events[self._last] or torch.cuda.Event()
-                self._events[self._last] = ev
-                ev.record(stream)
-            i = self._next
-            ev = self._events[i]
+            if self._streams[cur] is not None:
+                ev = self._events[cur] or torch.cuda.Event()
+                self._events[cur] = ev
+                ev.record(self._streams[cur])
+            ev = self._events[nxt]
             if ev is not None:
                 ev.synchronize()
-            self._streams[i] = stream
-        else:
-            i = self._next
-        self._next = (i + 1) % self._depth
-        self._last = i
-        out = self._bufs[i][:n]
+            self._streams[nxt] = torch.cuda.current_stream()
+        self._set = nxt
+        self._used = 0
+
+    def take(self, values) -> torch.Tensor:
+        n = len(values)
+        bufs = self._sets[self._set]
+        i = self._used
+        if i == len(bufs):
+            bufs.append(self._alloc(max(n, self._numel)))
+        elif n > bufs[i].numel():
+            bufs[i] = self._alloc(max(n, bufs[i].numel() * 2))
+        self._used = i + 1
+        out = bufs[i][:n]
         out.numpy()[:] = values
         return out
 

mstar/utils/knobs.py

@@ -49,7 +49,7 @@ def kv_chain_lazy_steps() -> int:
 
 
 def kv_plan_cache() -> int:
-    """``MSTAR_KV_PLAN_CACHE`` (default 1): the KV manager plans a leased
+    """``MSTAR_KV_PLAN_CACHE`` (default 0): the KV manager plans a leased
     decode step from 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: lengths advance by one, pages change only where a row crossed
@@ -61,12 +61,12 @@ def kv_plan_cache() -> int:
 
 
 def kv_pinned_indptrs() -> bool:
-    """``MSTAR_KV_PINNED_INDPTRS`` (default 1): the KV plan writes FlashInfer's
+    """``MSTAR_KV_PINNED_INDPTRS`` (default 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. The copies
     are then really asynchronous; from pageable memory the driver stages each
-    one, and FlashInfer makes the indices copy blocking. Off by default: the
-    event bookkeeping per take cost about what the staged copies did at
-    c32-c128 on 0.8B (measured neutral to slightly negative)."""
+    one, and FlashInfer makes the indices copy blocking. Off by default: no
+    measurable gain at c32-c128 on 0.8B. The cost is one CUDA event per plan,
+    which a buffer set waits on before it is rewritten."""
     return os.environ.get("MSTAR_KV_PINNED_INDPTRS", "0") == "1"

test/cpu-core.txt

@@ -79,5 +79,6 @@ test/modular/test_ingest_rollback.py
 test/modular/test_engine_uniform_inputs.py
 test/modular/test_kv_decode_plan_cache.py
 test/modular/test_kv_pinned_index_ring.py
+test/modular/test_kv_multi_label_plans.py
 test/modular/test_kv_plan_host_lens.py
 test/modular/test_sampler_trash_row.py