attn: skip a plan the wrapper already holds, in the ragged base

.claude/skills/engine-resources/SKILL.md

@@ -56,6 +56,7 @@ KV pages and recurrent (GDN/Mamba) state slots are two forms of the same per-req
 ## Ragged (cacheless) attention
 
 - **A ragged resource plans one layout per label.** A step declares, per request, the spans its layers attend over as `(label, length)` segments; the manager plans each label across the batch and a layer asks for its label's callable. Spans with different head geometries need separate resources, since a resource has one head layout.
+- **A wrapper skips a plan it already holds.** `_RaggedPrefillBase._plan` compares the layout with what that (label, slot)'s schedule describes, since `run` only reads it. A new wrapper on that base gets this for free; anything that invalidates a schedule without a plan has to clear it (`forget_plan`, as a dropped pre-plan does: its H2Ds ran on the plan stream and nothing orders the replacing step against them).
 - **Cross-attention between spans of a request is its own kind** (`RaggedCrossAttentionSpec`, #372): its step names `(q_label, kv_label)` pairs and plans each pair under `cross_label(q, kv)`. It keeps the self-attention resource's API unchanged; extending a resource that almost fits this way beats a model calling SDPA (see AGENTS.md invariant 1).
 
 ## Scope

mstar/engine/resources/attn/dense.py

@@ -195,8 +195,6 @@ def __init__(
 
         self._wrapper_class = wrapper_class
         self._wrappers: dict[EagerSlotKey, PrefixedPrefillWrapper] = {}
-        # wrapper -> the (cu_q, cu_k, causal) it was last planned with
-        self._planned: dict[EagerSlotKey, tuple] = {}
         self._wrapper_kwargs = dict(
             num_qo_heads=kv_config.num_qo_heads,
             num_kv_heads=kv_config.num_kv_heads,
@@ -331,23 +329,17 @@ def _build_plan(self, kv_out: KVPlanOutput, causal: bool, slot_key: EagerSlotKey
             max_q = max(max_q, view.to_compute)
             max_k = max(max_k, view.length)
 
-        # A wrapper holds the schedule of its last plan until it is planned
-        # again (``run`` only reads it), and a denoising request keeps its
-        # layout from step to step, so an unchanged layout is not replanned:
-        # FlashInfer's plan is a device kernel plus a stream sync.
+        # Fresh pinned host layouts per step: the FA3 wrapper H2Ds them non-blocking
+        # (from pageable memory that copy would sync the stream), FlashInfer's plan
+        # reads them on the host, and neither happens when the wrapper already
+        # holds this layout.
         wrapper = self._wrapper(slot_key.label, slot_key.slot)
-        layout = (tuple(cu_q), tuple(cu_k), causal)
-        if self._planned.get(slot_key) != layout:
-            # Fresh pinned host layouts per plan: the FA3 wrapper H2Ds them
-            # non-blocking (from pageable memory that copy would sync the
-            # stream), FlashInfer's plan reads them on the host.
-            pin = self._on_cuda
-            wrapper.plan(
-                torch.tensor(cu_q, dtype=torch.int32, pin_memory=pin),
-                torch.tensor(cu_k, dtype=torch.int32, pin_memory=pin),
-                causal=causal,
-            )
-            self._planned[slot_key] = layout
+        pin = self._on_cuda
+        wrapper.plan(
+            torch.tensor(cu_q, dtype=torch.int32, pin_memory=pin),
+            torch.tensor(cu_k, dtype=torch.int32, pin_memory=pin),
+            causal=causal,
+        )
         return DensePlan(
             segments=tuple(segments),
             cu_q=tuple(cu_q),

mstar/engine/resources/attn/ragged/base.py

@@ -355,6 +355,10 @@ def plan(self, step, ctx: StepContext):
         self._preplanned = ctx.is_preplan
 
     def clear_preplan(self):
+        # The dropped plan's H2Ds went on the plan stream, and the step that runs
+        # instead waits on no event from it, so it must stage its layout itself.
+        for wrapper in self._preplan_states.values():
+            wrapper.forget_plan()
         # rebind rather than clear: a consumed preplan dict is the live one
         self._preplanned = False
         self._preplan_states = {}
@@ -416,7 +420,6 @@ def plan(self, step: AttentionStep, ctx: StepContext):
     def _plan_layouts(self, step: AttentionStep, ctx: StepContext, plan_states) -> None:
         for label, segments in self._group_segments_by_label(step.segments or ()).items():
             wrapper = self._wrapper_for(ctx.slot_lease, ctx.slot, label)
-            # TODO: cache the latest plan state
             wrapper.plan(cu_seqlens=self._cu_seqlens(segments), causal=step.causal)
             plan_states[label] = wrapper
 

mstar/engine/resources/attn/ragged/block_causal.py

@@ -137,6 +137,9 @@ def __init__(
         # bucket's worst case so a small first plan can't cap it.
         self.plan([max_total_tokens])
 
+    def forget_plan(self) -> None:
+        """Every ``plan`` stages its layout here, so there is nothing to forget."""
+
     @property
     def num_segments(self) -> int:
         """Real (unpadded) segment count from the most recent ``plan``."""

mstar/engine/resources/attn/ragged/wrappers.py

@@ -25,7 +25,8 @@ def padded_head_dim(head_dim: int) -> int:
 
 class _RaggedPrefillBase(ABC):
     """What the self- and cross-attention wrappers share: the graph-mode static
-    buffers and their padding. A backend subclass plans and runs its kernel.
+    buffers and their padding, and skipping a plan the wrapper already holds. A
+    backend subclass plans and runs its kernel.
     """
 
     def __init__(
@@ -54,6 +55,8 @@ def __init__(
             max_total_kv_tokens if max_total_kv_tokens is not None else max_total_tokens
         )
         self._num_segments = 0
+        # the layout this wrapper's schedule describes; see ``_plan``
+        self._planned_layout: tuple | None = None
 
         if accelerator_graph:
             assert max_num_segments is not None, "max_num_segments required for accelerator graph mode"
@@ -117,13 +120,37 @@ def _prepare_cu_seqlens(
         return cu, total_tokens
 
     def _plan(self, cu_seqlens: torch.Tensor, kv_cu_seqlens: torch.Tensor, causal: bool) -> None:
+        # A wrapper keeps the schedule of its last plan until it is planned again
+        # and ``run`` only reads it, so a layout it already holds is not replanned:
+        # a plan costs host time, plus a stream sync on a backend whose plan reads
+        # the layout back from the device, and a denoising request keeps its
+        # layout from step to step. Each (label, slot) has its own wrapper, so this
+        # compares against what this slot's schedule says, never against what the
+        # previous step ran.
+        layout = self._layout_key(cu_seqlens, kv_cu_seqlens, causal)
+        if layout == self._planned_layout:
+            return
         cu, total_tokens = self._prepare_cu_seqlens(cu_seqlens, self.max_total_tokens)
         if kv_cu_seqlens is cu_seqlens:
             kv_cu = cu
         else:
             kv_cu, _ = self._prepare_cu_seqlens(kv_cu_seqlens, self.max_total_kv_tokens, "key ")
         self._num_segments = int(cu_seqlens.numel()) - 1
+        # a plan that raises may already have overwritten the schedule
+        self._planned_layout = None
         self._plan_kernel(cu, kv_cu, total_tokens, causal)
+        self._planned_layout = layout
+
+    @staticmethod
+    def _layout_key(cu_seqlens: torch.Tensor, kv_cu_seqlens: torch.Tensor, causal: bool) -> tuple:
+        q = cu_seqlens.tolist()
+        return q, q if kv_cu_seqlens is cu_seqlens else kv_cu_seqlens.tolist(), causal
+
+    def forget_plan(self) -> None:
+        """Replan on the next ``plan``, whatever its layout. For a plan whose
+        staging the wrapper's next user can't order against (a dropped pre-plan,
+        whose H2Ds went on the plan stream with nothing left to wait on them)."""
+        self._planned_layout = None
 
     @abstractmethod
     def _plan_kernel(

test/modular/test_dense_attention.py

@@ -140,38 +140,6 @@ def test_cumulative_lengths_pack_in_kv_plan_order(self):
         assert [s.request_id for s in plan.segments] == ["r0", "r1"]
         assert plan.causal is False
 
-    def test_an_unchanged_layout_is_not_replanned(self, monkeypatch):
-        """A wrapper keeps its last schedule, so a step whose layout matches
-        what that wrapper was last planned with skips the plan; another
-        layout, another causal flag or the other slot's wrapper plans again."""
-        calls = []
-        real_plan = DensePrefixedPrefillWrapper.plan
-
-        def counting_plan(self, *args, **kwargs):
-            calls.append(self)
-            return real_plan(self, *args, **kwargs)
-
-        monkeypatch.setattr(DensePrefixedPrefillWrapper, "plan", counting_plan)
-        manager = _manager()
-        views = [_view("r0", [0, 1], prefix=3, fresh=2)]
-        step = AttentionStep(causal=False)
-
-        manager.plan(step, _ctx(views))
-        manager.plan(step, _ctx(views))
-        assert len(calls) == 1
-        # the skipped plan still hands the layer the planned wrapper
-        assert manager._current_plans["main"].wrapper is calls[0]
-
-        manager.plan(step, _ctx(views, slot=1))
-        assert len(calls) == 2 and calls[1] is not calls[0]
-
-        manager.plan(AttentionStep(causal=True), _ctx(views))
-        assert len(calls) == 3
-        manager.plan(step, _ctx([_view("r0", [0, 1], prefix=3, fresh=3)]))
-        assert len(calls) == 4
-        manager.plan(step, _ctx([_view("r0", [0, 1], prefix=3, fresh=3)]))
-        assert len(calls) == 4
-
     def test_causal_rides_on_the_step(self):
         manager = _manager()
         manager.plan(AttentionStep(causal=True), _ctx([_view("r0", [0], prefix=1, fresh=1)]))