feat(exact-prefix): add direct CUDA transfers for Fluxon Plan API

cacheseek/reuse/exact_prefix/__init__.py

@@ -19,7 +19,7 @@
   - the version field must hash the real config blob (keys.config_blob_hash).
 """
 
-from cacheseek.stores.tier import InMemoryTierStore, TensorStoreTierStore
+from cacheseek.stores.tier import InMemoryTierStore, PlanTensorTierStore, TensorStoreTierStore
 
 from .config import ModelGeometry, WorldKVConfig, bytes_per_chunk_kv, calibrate_break_even_k
 from .keys import build_action_chain, config_blob_hash, derive_seed, node_key, root_hash
@@ -39,6 +39,7 @@
     "save_forest_snapshot", "load_forest_snapshot",
     "WorldKVManager", "WorldKVConfig", "ModelGeometry", "FastForwardResult",
     "KVTierStore", "RollingWindow", "InMemoryTierStore", "TensorStoreTierStore",
+    "PlanTensorTierStore",
     "build_action_chain", "config_blob_hash", "derive_seed", "node_key", "root_hash",
     "bytes_per_chunk_kv", "calibrate_break_even_k",
 ]

cacheseek/reuse/exact_prefix/manager.py

@@ -149,6 +149,14 @@ def materialize(self, node: TrieNode, window: RollingWindow) -> bool:
         for n in path:
             n.ref_count += 1                         # materialize in flight; eviction must not reclaim
         try:
+            materialize_path = getattr(self.store, "materialize_path", None)
+            if callable(materialize_path):
+                ok = bool(materialize_path(path, window, depth=node.depth))
+                if ok:
+                    now = self._now()
+                    for n in path:
+                        n.last_access = now
+                return ok
             n_layers = path[-1].blob.n_layers        # type: ignore[union-attr]
             for layer in range(n_layers):
                 blobs = [(n.depth, self.store.get_layer(n.blob, layer)) for n in path]

cacheseek/reuse/exact_prefix/telefuser_lingbot.py

@@ -41,6 +41,7 @@
 import torch
 
 from cacheseek.service.query import CacheQuery
+from cacheseek.stores.cuda_transfer import PointerCopy
 
 from .config import WorldKVConfig
 from .keys import canonical_json_bytes, config_blob_hash, root_hash, sha256
@@ -86,9 +87,15 @@ def make_rolling_config(
     "max_attention_size",
     "max_sequence_length",
 )
+FLUXON_PLAN_LAYOUT_VERSION = "fluxon-plan-python-cuda-v1"
 
 
-def session_root_hash(session_config: Any, *, model_fingerprint: bytes) -> bytes:
+def session_root_hash(
+    session_config: Any,
+    *,
+    model_fingerprint: bytes,
+    runtime: Any | None = None,
+) -> bytes:
     """Compute the namespace root_hash for a TeleFuser session.
 
     Combines an image fingerprint (mode, size, raw pixel bytes), a normalized
@@ -113,6 +120,18 @@ def session_root_hash(session_config: Any, *, model_fingerprint: bytes) -> bytes
     )
     prompt_fp = sha256(b"prompt", session_config.prompt.strip().encode("utf-8"))
     blob = {f: getattr(session_config, f, None) for f in SESSION_KEY_FIELDS}
+    if runtime is not None:
+        kv_tensor = runtime.self_kv_cache[0]["k"]
+        blob["fluxon_plan_layout"] = {
+            "storage_format": FLUXON_PLAN_LAYOUT_VERSION,
+            "n_layers": len(runtime.self_kv_cache),
+            "kv_shape": list(kv_tensor.shape),
+            "kv_dtype": str(kv_tensor.dtype),
+            "latent_shape": list(runtime.noise_chunks[0].shape),
+            "latent_dtype": str(kv_tensor.dtype),
+            "chunk_size": int(runtime.chunk_size),
+            "frame_tokens": int(runtime.frame_tokens),
+        }
     cfg_hash = config_blob_hash(blob, weights_fingerprint=model_fingerprint)
     return root_hash(image_fp=image_fp, prompt_fp=prompt_fp, config_blob_hash=cfg_hash)
 
@@ -172,6 +191,7 @@ class _RingKVWindow:
     def __init__(self, runtime: Any) -> None:
         self._rt = runtime
         self._local_end_tokens = 0
+        self._planned_local_end_tokens = 0
 
     def _frames_to_seed(self, layer_kv: dict, depth: int) -> list[int]:
         rt = self._rt
@@ -216,14 +236,104 @@ def seed_layer(self, layer: int, blobs: list[tuple[int, Any]], depth: int) -> No
             )
         self._local_end_tokens = len(frames) * ft
 
+    def chunk_value_nbytes(self, layer: int) -> int:
+        """Return the raw byte length of one chunk's K or V value."""
+        kv = self._rt.self_kv_cache[layer]
+        batch, _, heads, head_dim = kv["k"].shape
+        return (
+            int(batch)
+            * int(self._rt.chunk_size)
+            * int(self._rt.frame_tokens)
+            * int(heads)
+            * int(head_dim)
+            * int(kv["k"].element_size())
+        )
+
+    @property
+    def cuda_device(self) -> Any:
+        return self._rt.self_kv_cache[0]["k"].device
+
+    def build_seed_copies(
+        self,
+        layer: int,
+        blobs: list[tuple[int, tuple[int, int]]],
+        depth: int,
+    ) -> list[PointerCopy]:
+        """Map chunk Plan pointers into contiguous runs in the physical KV ring."""
+        rt = self._rt
+        kv = rt.self_kv_cache[layer]
+        k_tensor, v_tensor = kv["k"], kv["v"]
+        if not k_tensor.is_cuda or not v_tensor.is_cuda:
+            raise ValueError("Fluxon Plan restore requires CUDA KV tensors")
+        if not k_tensor.is_contiguous() or not v_tensor.is_contiguous():
+            raise ValueError("Fluxon Plan restore requires contiguous KV tensors")
+        if k_tensor.shape != v_tensor.shape or k_tensor.dtype != v_tensor.dtype:
+            raise ValueError("runtime K/V layout mismatch")
+
+        frame_tokens = int(rt.frame_tokens)
+        chunk_frames = int(rt.chunk_size)
+        batch, buffer_tokens, heads, head_dim = map(int, k_tensor.shape)
+        token_bytes = heads * head_dim * int(k_tensor.element_size())
+        source_row_bytes = chunk_frames * frame_tokens * token_bytes
+        destination_row_bytes = buffer_tokens * token_bytes
+        by_depth = dict(blobs)
+        frames = self._frames_to_seed(kv, depth)
+
+        runs: list[tuple[int, int, int, int]] = []
+        for position, global_frame in enumerate(frames):
+            source_chunk = global_frame // chunk_frames
+            source_frame = global_frame % chunk_frames
+            if source_chunk not in by_depth:
+                raise KeyError(f"missing source chunk {source_chunk}")
+            if runs:
+                old_chunk, old_position, old_source_frame, count = runs[-1]
+                if (
+                    old_chunk == source_chunk
+                    and old_position + count == position
+                    and old_source_frame + count == source_frame
+                ):
+                    runs[-1] = (old_chunk, old_position, old_source_frame, count + 1)
+                    continue
+            runs.append((source_chunk, position, source_frame, 1))
+
+        copies: list[PointerCopy] = []
+        for source_chunk, destination_frame, source_frame, frame_count in runs:
+            source_k, source_v = by_depth[source_chunk]
+            copy_bytes = frame_count * frame_tokens * token_bytes
+            for batch_index in range(batch):
+                source_offset = (
+                    batch_index * source_row_bytes + source_frame * frame_tokens * token_bytes
+                )
+                destination_offset = (
+                    batch_index * destination_row_bytes
+                    + destination_frame * frame_tokens * token_bytes
+                )
+                copies.extend(
+                    [
+                        PointerCopy(
+                            dst=int(k_tensor.data_ptr()) + destination_offset,
+                            src=int(source_k) + source_offset,
+                            nbytes=copy_bytes,
+                        ),
+                        PointerCopy(
+                            dst=int(v_tensor.data_ptr()) + destination_offset,
+                            src=int(source_v) + source_offset,
+                            nbytes=copy_bytes,
+                        ),
+                    ]
+                )
+        self._planned_local_end_tokens = len(frames) * frame_tokens
+        return copies
+
     def set_resume_depth(self, depth: int) -> None:
         """Set each layer's global_end_index (logical, F*ft) and local_end_index
         (physical buffer fill from the last seed) so the DiT resumes at chunk depth+1."""
         rt = self._rt
         global_end = (depth + 1) * rt.chunk_size * rt.frame_tokens
+        local_end = self._planned_local_end_tokens or self._local_end_tokens
         for kv in rt.self_kv_cache:
             kv["global_end_index"] = global_end
-            kv["local_end_index"] = self._local_end_tokens
+            kv["local_end_index"] = local_end
 
 
 class LingBotWorldKVBinding:
@@ -262,8 +372,8 @@ def __init__(
         self.ingest_enabled = ingest_enabled
         # Optional cross-process hits: if the forest is empty at startup, load the
         # index from a snapshot; write it back between sessions / after finalize.
-        # Only meaningful with a persistent store (TensorStoreTierStore over
-        # LocalDisk/Fluxon); InMemory lives only within the process. See
+        # Only meaningful with a persistent store (the generic LocalDisk adapter
+        # or the synchronous Fluxon Plan adapter); InMemory is process-local. See
         # docs/design_exact_prefix_reuse/04-physical-view.md.
         self.snapshot_path = snapshot_path
         self.snapshot_on_finalize = snapshot_on_finalize
@@ -298,7 +408,9 @@ def on_runtime_created(self, runtime: Any, session_config: Any) -> None:
         from .keys import build_action_chain
 
         root = session_root_hash(
-            session_config, model_fingerprint=self.model_fingerprint
+            session_config,
+            model_fingerprint=self.model_fingerprint,
+            runtime=runtime,
         )
         self._ns = self.forest.get_or_create_namespace(root, root)
         self._actions = chunk_action_keys(runtime)
@@ -316,33 +428,56 @@ def on_runtime_created(self, runtime: Any, session_config: Any) -> None:
         # break-even gate); materialization/latent/RNG are engine-adapter
         # responsibilities (interpreting the FastForward hint) and stay in this binding.
         res = asyncio.run(self.strategy.lookup(self._query))
+        self._parent = self._ns.root
+        self.last_fast_forward = 0
+        runtime.world_kv_cached_latents = {}
         if not res.hit:
-            self._parent = self._ns.root
-            self.last_fast_forward = 0
             return
         hint = res.resume_hint
+        path_from_hit: list[TrieNode] = []
+        node = hint.node
+        while node is not None and node.depth >= 0:
+            path_from_hit.append(node)
+            node = node.parent
+        nodes = list(reversed(path_from_hit))
+        if not nodes or any(node.skeleton is None for node in nodes):
+            return
+
+        cached: dict[int, torch.Tensor] = {}
+        materialize_skeletons = getattr(self.mgr.store, "materialize_skeletons", None)
+        if callable(materialize_skeletons):
+            latent_dtype = runtime.self_kv_cache[0]["k"].dtype
+            latent_device = runtime.self_kv_cache[0]["k"].device
+            targets = {
+                node.depth: torch.empty(
+                    runtime.noise_chunks[node.depth].shape,
+                    dtype=latent_dtype,
+                    device=latent_device,
+                )
+                for node in nodes
+            }
+            target_items = [
+                (node.skeleton.latent_locator, targets[node.depth]) for node in nodes
+            ]
+            if not materialize_skeletons(target_items):
+                return
+            cached.update(targets)
+        else:
+            for node in nodes:
+                latent = self.mgr.store.get_skeleton(node.skeleton.latent_locator)
+                if latent is None:
+                    return
+                cached[node.depth] = latent
+
         if not self.mgr.materialize(hint.node, _RingKVWindow(runtime)):
-            self._parent = self._ns.root  # incomplete window -> fall back to cold run
-            self.last_fast_forward = 0
             return
+
         self._parent = hint.node
         k = hint.k
         self.last_fast_forward = k
-
-        # 1. Skipped chunks -> decode-only: take the latent from the skeleton, no
-        #    denoise / no rewrite.
-        cached: dict[int, torch.Tensor] = {}
-        path: list[TrieNode] = []
-        n = hint.node
-        while n is not None and n.depth >= 0:
-            path.append(n)
-            n = n.parent
-        for node in reversed(path):  # chunk 0..K-1
-            latent = self.mgr.store.get_skeleton(node.skeleton.latent_locator)
-            cached[node.depth] = latent
         runtime.world_kv_cached_latents = cached
 
-        # 2. Burn the generator draws for skipped chunks (len(timesteps)-1 per
+        # Burn the generator draws for skipped chunks (len(timesteps)-1 per
         #    chunk, shape=chunk latent, dtype=bf16, one-to-one with denoise_chunk's
         #    torch.randn). Without this, the RNG stream is misaligned from chunk K
         #    onward and exact replay silently breaks.
@@ -370,6 +505,19 @@ def on_chunk_finalized(
         latent and ingest them."""
         if not self.ingest_enabled or self._ns is None:
             return
+        expected_shape = tuple(runtime.noise_chunks[idx].shape)
+        actual_shape = tuple(denoised.shape)
+        if actual_shape != expected_shape:
+            raise ValueError(
+                f"latent shape mismatch: expected={expected_shape} got={actual_shape}"
+            )
+        expected_dtype = runtime.self_kv_cache[0]["k"].dtype
+        if denoised.dtype != expected_dtype:
+            raise ValueError(
+                f"latent dtype mismatch: expected={expected_dtype} got={denoised.dtype}"
+            )
+
+        plan_data_path = callable(getattr(self.mgr.store, "materialize_skeletons", None))
         ct = runtime.chunk_size * runtime.frame_tokens
         payload = []
         for kv in runtime.self_kv_cache:
@@ -378,13 +526,19 @@ def on_chunk_finalized(
             # full-length mode); local_end was just advanced by the clean rewrite.
             e = int(kv["local_end_index"])
             s = e - ct
-            payload.append(
-                (
-                    kv["k"][:, s:e].detach().to("cpu").clone(),
-                    kv["v"][:, s:e].detach().to("cpu").clone(),
-                )
-            )
-        latent = denoised.detach().to("cpu").clone()
+            k_view = kv["k"][:, s:e].detach()
+            v_view = kv["v"][:, s:e].detach()
+            if not plan_data_path:
+                # Generic stores may enqueue writes, so force caller-owned CPU
+                # snapshots before the runtime reuses or overwrites these ring slots.
+                k_view = k_view.to(device="cpu", copy=True)
+                v_view = v_view.to(device="cpu", copy=True)
+            payload.append((k_view, v_view))
+        if plan_data_path:
+            latent = denoised.detach()
+        else:
+            # The generic skeleton payload needs the same async-lifetime isolation.
+            latent = denoised.detach().to(device="cpu", copy=True)
         # Writeback goes through the shared Strategy protocol; chunk data is passed
         # via ctx (exact save is chunk-granular streaming).
         ctx = {

cacheseek/stores/__init__.py

@@ -4,9 +4,9 @@
 
 Capability tiers:
   - Bytes contract: ``KVStore`` (put/get/remove/list_keys over bytes).
-  - Tensor contract: ``TensorKVStore`` (put_tensor/get_tensor, optional zero-copy).
-  - Tier adapter: ``TensorStoreTierStore`` (spec bookkeeping + async write queue
-    + per-layer (k, v) splitting).
+  - Tensor contract: ``TensorKVStore`` (optional generic tensor adapters).
+  - Tier adapters: ``TensorStoreTierStore`` (per-key tensor get/put) and
+    ``PlanTensorTierStore`` (Fluxon exact-prefix Plan pointer capability).
 Backends: memory / local_file / fluxon (the bytes trio) plus InMemoryTierStore /
 LocalDiskTensorStore.
 
@@ -23,6 +23,7 @@
     "KVStore", "TensorKVStore", "Tier", "BlobHandle",
     "InMemoryKVStore", "LocalFileKVStore", "FluxonKVStore",
     "InMemoryTierStore", "LocalDiskTensorStore", "TensorStoreTierStore",
+    "PlanTensorTierStore",
 ]
 
 _LAZY: dict[str, tuple[str, str]] = {
@@ -32,6 +33,7 @@
     "InMemoryTierStore": ("cacheseek.stores.tier", "InMemoryTierStore"),
     "LocalDiskTensorStore": ("cacheseek.stores.tier", "LocalDiskTensorStore"),
     "TensorStoreTierStore": ("cacheseek.stores.tier", "TensorStoreTierStore"),
+    "PlanTensorTierStore": ("cacheseek.stores.tier", "PlanTensorTierStore"),
 }
 
 

cacheseek/stores/base.py

@@ -3,8 +3,8 @@
 """KVStore Protocol — opaque byte-blob storage keyed by string id.
 
 ``TensorKVStore`` is an OPTIONAL capability layered on top: backends that can
-store/return tensors without serializing to bytes (e.g. Fluxon via DLPack)
-implement it. Callers route through ``adapters/lingbot_fast/tensor_block_io``,
+store/return tensors without serializing to bytes implement it. Callers route
+through ``adapters/lingbot_fast/tensor_block_io``,
 which falls back to a pickle-free raw-bytes path on stores that only satisfy
 ``KVStore``. ``KVStore`` itself is unchanged — the capability is additive and
 non-breaking; existing ``put(bytes)`` / ``get`` callers are untouched.
@@ -74,10 +74,10 @@ def list_keys(self) -> list[str]:
 class TensorKVStore(Protocol):
     """Optional zero-copy tensor capability on top of ``KVStore``.
 
-    A backend implementing this can ingest/return torch tensors directly
-    (Fluxon hands the DLPack pointer to its Rust layer — no Python bytes, no
-    pickle). ``isinstance(store, TensorKVStore)`` is the routing check; stores
-    without these methods fall back to a pickle-free raw-bytes path.
+    A backend implementing this can ingest/return torch tensors directly.
+    ``isinstance(store, TensorKVStore)`` is the routing check; stores without
+    these methods fall back to a pickle-free raw-bytes path. Fluxon exact-prefix
+    reuse instead uses its explicit Plan pointer capability.
 
     Contract:
     - ``shape`` / ``dtype`` on ``get_tensor`` are REQUIRED — raw bytes are not