Offer the host store only checkpoints a request may resume

Every published checkpoint was copied to the host store at publish time.
On a growing conversation about four checkpoints are published per turn
and the next request resumes at one: a message checkpoint is shadowed a
few tokens later by its own prompt checkpoint, and a reply-end
checkpoint is rarely resumed by clients that re-render the reply.

The scheduler now defers the host offer (never the GPU capture) of a
message checkpoint whose own prompt checkpoint follows less than one
page later, and of terminal checkpoints. A deferred checkpoint is
offered when a request resumes exactly at it (it is still on the GPU and
pinned by that request), or, for a message checkpoint whose prompt
capture could not get a snapshot, at that point while the request still
pins it. A preempted request's terminal checkpoint counts as resumed, and a
deferred message checkpoint is offered when its request is cancelled or
fails.
The parent reported for a publication is the nearest offered ancestor,
with the tail pages since it. When the GPU compacts an offered
checkpoint whose heir was not offered, the pool hears the first offered
descendant as the heir, so announced chains stay connected;
a deferred checkpoint that became a branch point is offered as the heir. Captures,
prefill cuts, GPU retention and eviction order are unchanged; the rule
uses only page size and the selected message boundaries.

Validation: upstream CI's Linux commands: selene with test-support 389
passed, selene-manager all passed, helios the same 1187 passed / 7
failed as main (main's DeepSeek tokenizer-fixture defect), scheduler
tests all pass. Tests cover deferral at page boundaries, offer on reuse,
the capture-failure and preemption paths, parent and tail reporting,
compaction with a deferred heir in pool mode, and output lengths 1 and
3. Independent review; findings addressed.
Not yet run on a GPU.

Design-Signed-Off-By: yatesdr; design-review=1c01

Author: yatesdr

docs/prefix-cache-offload.md

@@ -28,8 +28,7 @@ defines configuration precedence.
 
 ## 1. Scope and configurations
 
-- Published prefix checkpoints at instruction, prompt, and terminal endpoints
-  are copied to the first host tier when published in both local and pool modes.
+- Offered GPU checkpoints are copied to the first host tier in local and pool modes.
   A later matching request restores a retained checkpoint the GPU no longer holds.
 - Live request state, in-flight speculative rounds, and anything not published
   as a checkpoint are never offloaded.
@@ -171,18 +170,26 @@ pub trait HostStore {
 }
 ```
 
-The worker calls `store` for each newly published checkpoint, without waiting for
-GPU pressure. `Queued` holds one GPU pin until `Stored` or `Deferred`; `Present`
+The worker calls `store` for each offered checkpoint, without waiting for GPU
+pressure. The scheduler defers a message checkpoint whose own planned prompt
+checkpoint is its direct child less than one physical page later, and terminal
+checkpoints, until a request resumes at that exact checkpoint. A preempted
+terminal counts as resumed. If its prompt capture cannot obtain a snapshot, the
+message is offered while the request still pins it; cancellation or error before
+that capture offers it at retirement. GPU captures, prefill cuts, and retention
+are unchanged; deferred points can be lost to GPU
+eviction before reuse, and are not announced to the pool until offered.
+`Queued` holds one GPU pin until `Stored` or `Deferred`; `Present`
 requires a committed local copy. `Retry` accepts no copy. The worker retains only
 checkpoint IDs for retries after store progress and discards IDs that have left
 the GPU cache. `Deferred` releases the copy pin and schedules the same retry when
 a queued publication cannot obtain host capacity. A failed transfer releases its
 pin without retrying. GPU eviction uses the existing LRU, leaf and compaction
 rules. `fetch` copies a looked-up chain into the ticket's page IDs and state slot.
 Submissions never wait; `poll` drains completions and failures.
-A `PublishedCheckpoint` names the GPU checkpoint, its key, its parent's key,
-its bucket digest, every page of
-its chain from the root with the count that belongs to the checkpoint itself,
+A `PublishedCheckpoint` names the GPU checkpoint, its key, its nearest offered
+ancestor's key, its bucket digest, every page of
+its chain from the root with the tail count since that offered ancestor,
 and the state slot of its row. `MemoryStore`, `DiskStore`, and `CompositeStore`
 implement the trait for serving, and a test double implements it for worker tests. The
 worker holds `Option<Box<dyn HostStore>>`.
@@ -566,12 +573,10 @@ API on this class of host and is the reference for call patterns.
   caller for cancellation and re-admission.
 - Cancelling a waiter removes it from the ticket. The ticket completes
   regardless, because the restored checkpoint is useful on its own.
-- `take_publications()` returns every checkpoint that became GPU-ready since
-  the previous call, each once, with its id, key, parent key, bucket digest,
-  the page ids of its whole chain with the count belonging to the checkpoint,
-  and the row's state slot from `snapshot_slot`. Checkpoints created by
-  `publish_boundary`, `retire`, and `finish_inbound` enter the list when the
-  completed fence makes them ready.
+- `take_publications()` returns GPU-ready checkpoints offered since the previous
+  call, each once, with their id, key, nearest offered ancestor's key, bucket
+  digest, whole page chain, tail count since that ancestor, and row state slot.
+  Deferred checkpoints enter the list on reuse or the fallback above.
 - `SchedulerStats` gains `host_lookups`, `host_hits`, `host_reused_tokens`,
   and `inbound_waiting`.
 - `release_after_shutdown` drops tickets.
@@ -629,7 +634,7 @@ the store and its transfer threads.
 
 Policy:
 
-- Outbound: copy each published checkpoint into the first host tier. In local
+- Outbound: copy each offered checkpoint into the first host tier. In local
   composite mode, spill to disk; in pool mode, move records down the
   eviction cascade only when the preceding tier needs room. Only queued copies
   pin GPU checkpoints. A capacity refusal waits for host progress before retry;
@@ -1008,7 +1013,7 @@ same completion events the store produces.
   Local mode uses the same parked/spill path and LRU unlink at the end of its
   chain, with the same completed-maintenance budget accounting. Only end-of-chain eviction,
   expiry, infeasible footprints (`Oversize`) and terminal disk I/O errors discard data.
-  GPU publication announces the GPU holding and starts its first host copy;
+  A host offer announces the GPU holding and starts its first host copy;
   GPU eviction sends the existing notice and never starts a late copy.
 - **Retention metrics.** `flare_host_cache_held_total{tier}` counts record
   holding events. `flare_host_cache_dropped_total{tier,reason}` counts local

helios/src/cache.rs

@@ -224,6 +224,16 @@ impl PrefixCheckpointCache {
         })
     }
 
+    pub(crate) fn single_child(&self, id: CheckpointId) -> Option<CheckpointId> {
+        if self.nodes.get(&id)?.child_count != 1 {
+            return None;
+        }
+        self.nodes
+            .values()
+            .find(|node| node.parent == Some(id))
+            .map(|node| node.id)
+    }
+
     /// The GPU-ready checkpoint published under `prefix`, if any.
     pub fn ready_by_prefix(&self, prefix: &PrefixKey) -> Option<CheckpointId> {
         let id = *self.prefixes.get(prefix)?;

helios/src/scheduler.rs

@@ -283,6 +283,10 @@ pub struct LiveRequest {
 }
 
 impl LiveRequest {
+    fn plans_prompt_capture(&self) -> bool {
+        self.checkpoint_budget >= 2 && self.max_new_tokens > 1
+    }
+
     fn prefill_kind(&self) -> GraphKind {
         if self.prompt_logprobs { GraphKind::ScoredPrefill } else { GraphKind::Prefill }
     }
@@ -681,6 +685,7 @@ pub struct RequestScheduler {
     tickets: HashMap<PrefixKey, Ticket>,
     track_publications: bool,
     publications: Vec<CheckpointId>,
+    offered_publications: BTreeSet<CheckpointId>,
     track_evictions: bool,
     evictions: Vec<EvictedCheckpoint>,
     evictions_lost: bool,
@@ -757,6 +762,7 @@ impl RequestScheduler {
             tickets: HashMap::new(),
             track_publications: false,
             publications: Vec::new(),
+            offered_publications: BTreeSet::new(),
             track_evictions: false,
             evictions: Vec::new(),
             evictions_lost: false,
@@ -1627,6 +1633,11 @@ impl RequestScheduler {
         self.stats.cache_lookups += u64::from(!spec.prompt_logprobs);
         self.stats.cache_hits += u64::from(hit);
         self.stats.reused_tokens += reused;
+        if inbound.is_none()
+            && let Some(id) = lookup.checkpoint
+        {
+            self.record_publication(id);
+        }
         Ok(Admitted { handle, ticket })
     }
 
@@ -1686,8 +1697,8 @@ impl RequestScheduler {
         }
         if matches!(result, PublishResult::Created { .. }) {
             self.stats.snapshot_publications += 1;
-            self.record_publication(id);
         }
+        self.record_publication(id);
         let snapshot = self.cache.checkpoint(id)?.snapshot();
         let shared_pages = self.cache.page_chain(id)?;
         for handle in ticket.waiters {
@@ -1720,8 +1731,8 @@ impl RequestScheduler {
         self.track_publications = true;
     }
 
-    /// Checkpoints that became GPU-ready since tracking was enabled or the
-    /// previous call, in that order and each once, omitting retired checkpoints.
+    /// GPU-ready checkpoints offered since the previous call, each once and
+    /// omitting retired checkpoints. Deferred checkpoints are offered on reuse.
     pub fn take_publications(&mut self) -> Vec<PublishedCheckpoint> {
         std::mem::take(&mut self.publications)
             .into_iter()
@@ -1731,23 +1742,32 @@ impl RequestScheduler {
 
     pub(crate) fn publication(&self, id: CheckpointId) -> Option<PublishedCheckpoint> {
         let checkpoint = self.cache.checkpoint(id).ok()?;
-        let parent = match checkpoint.parent() {
-            Some(parent) => Some(self.cache.checkpoint(parent).ok()?.prefix()),
-            None => None,
-        };
+        let mut ancestor = checkpoint.parent();
+        let mut parent = None;
+        while let Some(ancestor_id) = ancestor {
+            let node = self.cache.checkpoint(ancestor_id).ok()?;
+            if self.offered_publications.contains(&ancestor_id) {
+                parent = Some(node.prefix());
+                break;
+            }
+            ancestor = node.parent();
+        }
+        let page_tokens = self.capabilities.page_tokens as u64;
         Some(PublishedCheckpoint {
             id,
             key: checkpoint.prefix(),
             parent,
             bucket: checkpoint.bucket(),
             chain_pages: self.cache.page_chain(id).ok()?,
-            tail_pages: checkpoint.tail_pages().len(),
+            tail_pages: (checkpoint.prefix().consumed_tokens / page_tokens
+                - parent.map_or(0, |key| key.consumed_tokens / page_tokens))
+                as usize,
             slot: self.snapshot_slot(checkpoint.snapshot()),
         })
     }
 
     fn record_publication(&mut self, id: CheckpointId) {
-        if !self.track_publications {
+        if !self.track_publications || !self.offered_publications.insert(id) {
             return;
         }
         // Bound queued notifications when the store delays draining them.
@@ -1760,6 +1780,7 @@ impl RequestScheduler {
 
     fn apply_completion(&mut self, update: CompletionUpdate) {
         self.stats.snapshot_publications += update.published.len() as u64;
+        // Boundary and terminal captures publish already ready, so deferral is decided there.
         for id in update.published {
             self.record_publication(id);
         }
@@ -2547,9 +2568,8 @@ impl RequestScheduler {
                 .or_else(|| {
                     (endpoint == request.prompt_tokens as u64
                         && endpoint > request.prefix_tokens
-                        && request.checkpoint_budget >= 2
-                        && request.max_new_tokens > 1)
-                        .then_some(BoundaryKind::Prompt)
+                        && request.plans_prompt_capture())
+                    .then_some(BoundaryKind::Prompt)
                 });
                 if let Some(kind) = capture {
                     captures.push((lane.request, Some(kind), endpoint, request.phase));
@@ -2606,6 +2626,12 @@ impl RequestScheduler {
             } else {
                 self.retain_snapshot()?
             };
+            if snapshot.is_none()
+                && matches!(kind, Some(BoundaryKind::Prompt))
+                && let Some(prefix) = self.request(handle)?.prefix
+            {
+                self.record_publication(prefix);
+            }
             let request = self.request_mut(handle)?;
             if let Some(snapshot) = snapshot {
                 if let Some(kind) = kind {
@@ -2656,12 +2682,20 @@ impl RequestScheduler {
             ready_after: request.last_use,
         };
         let history = request.history.clone();
+        // Only the last selected message can be the prompt's direct parent.
+        let defer = matches!(capture.kind, BoundaryKind::Message)
+            && request.message_boundaries.last() == Some(&(capture.consumed_tokens as usize))
+            && request.prompt_tokens - (capture.consumed_tokens as usize)
+                < self.capabilities.page_tokens
+            && request.plans_prompt_capture();
         let result = self.cache.publish(&history, &proposal)?;
         if let PublishResult::Created { id, ready } = result {
             if !ready {
                 return Err("completed boundary copy was not GPU-ready".into());
             }
-            self.record_publication(id);
+            if !defer {
+                self.record_publication(id);
+            }
             self.cache.pin(id)?;
             if let Some(parent) = self.request(handle)?.prefix {
                 self.cache.unpin(parent)?;
@@ -2790,18 +2824,28 @@ impl RequestScheduler {
         let mut request = self.slots[handle.slot as usize].request.take().unwrap();
         let mut snapshot = request.terminal_snapshot;
         if matches!(result, Some(PublishResult::Created { .. })) {
-            if let Some(PublishResult::Created { id, ready: true }) = result {
+            if let Some(PublishResult::Created { ready: true, .. }) = result {
                 self.stats.snapshot_publications += 1;
-                self.record_publication(id);
             }
             request
                 .private_pages
                 .drain(..publication.as_ref().unwrap().tail_pages.len());
             snapshot = None;
         }
         if let Some(prefix) = request.prefix {
+            if !publish || request.error_code.is_some() {
+                self.record_publication(prefix);
+            }
             self.cache.unpin(prefix)?;
         }
+        if request.preempted
+            && let Some(
+                PublishResult::Created { id, ready: true }
+                | PublishResult::Existing { id, ready: true },
+            ) = result
+        {
+            self.record_publication(id);
+        }
         self.slots[handle.slot as usize].retiring = true;
         request.private_pages.extend(request.reserved_pages);
         self.retired.push(RetiredRequest {