diff --git a/cacheseek/reuse/exact_prefix/__init__.py b/cacheseek/reuse/exact_prefix/__init__.py index 69e411e..d4c9be8 100644 --- a/cacheseek/reuse/exact_prefix/__init__.py +++ b/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", ] diff --git a/cacheseek/reuse/exact_prefix/manager.py b/cacheseek/reuse/exact_prefix/manager.py index f15ab50..cc3d0c6 100644 --- a/cacheseek/reuse/exact_prefix/manager.py +++ b/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] diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index f66553f..d852500 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/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 = { diff --git a/cacheseek/stores/__init__.py b/cacheseek/stores/__init__.py index 35d6812..787110d 100644 --- a/cacheseek/stores/__init__.py +++ b/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"), } diff --git a/cacheseek/stores/base.py b/cacheseek/stores/base.py index 03bee65..7a76139 100644 --- a/cacheseek/stores/base.py +++ b/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 diff --git a/cacheseek/stores/cuda_transfer.py b/cacheseek/stores/cuda_transfer.py new file mode 100644 index 0000000..f5488d6 --- /dev/null +++ b/cacheseek/stores/cuda_transfer.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import threading +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True, slots=True) +class PointerCopy: + dst: int + src: int + nbytes: int + + +class CudaTransferRuntime: + """Small injectable facade around the CUDA runtime transfer APIs.""" + + _registered: set[tuple[int, int, int]] = set() + _register_lock = threading.Lock() + + def __init__(self, *, api: Any | None = None) -> None: + if api is None: + try: + import cuda.bindings.runtime as api + except ImportError as exc: + raise ImportError( + "Fluxon direct GPU transfer requires cuda-bindings; install cacheseek[fluxon]" + ) from exc + self._api = api + + @staticmethod + def _code(result: Any) -> int: + value = result[0] if isinstance(result, tuple) else result + return int(value) + + def _enum(self, enum_name: str, member: str) -> Any: + direct = getattr(self._api, member, None) + if direct is not None: + return direct + return getattr(getattr(self._api, enum_name), member) + + def _check(self, result: Any, op: str) -> None: + code = self._code(result) + success = int(self._enum("cudaError_t", "cudaSuccess")) + if code != success: + raise RuntimeError(f"{op} failed with CUDA error {code}") + + def _copy_batch(self, copies: Iterable[PointerCopy], *, kind: Any, stream: int) -> None: + materialized = tuple(copies) + if not materialized: + return + for index, copy in enumerate(materialized): + if copy.dst <= 0 or copy.src <= 0 or copy.nbytes <= 0: + raise ValueError(f"invalid pointer copy at index {index}: {copy!r}") + submitted = 0 + try: + for index, copy in enumerate(materialized): + self._check( + self._api.cudaMemcpyAsync(copy.dst, copy.src, copy.nbytes, kind, stream), + f"cudaMemcpyAsync[{index}]", + ) + submitted += 1 + except Exception as copy_error: + # The caller may release or abort Plan-owned host memory as soon as + # this method raises, so drain any copies that were already queued. + if submitted: + try: + self._check( + self._api.cudaStreamSynchronize(stream), + "cudaStreamSynchronize after cudaMemcpyAsync failure", + ) + except Exception as sync_error: + raise RuntimeError( + f"{copy_error}; cleanup synchronization also failed: {sync_error}" + ) from copy_error + raise + self._check(self._api.cudaStreamSynchronize(stream), "cudaStreamSynchronize") + + def copy_h2d(self, copies: Iterable[PointerCopy], *, stream: int) -> None: + kind = self._enum("cudaMemcpyKind", "cudaMemcpyHostToDevice") + self._copy_batch(copies, kind=kind, stream=stream) + + def copy_d2h(self, copies: Iterable[PointerCopy], *, stream: int) -> None: + kind = self._enum("cudaMemcpyKind", "cudaMemcpyDeviceToHost") + self._copy_batch(copies, kind=kind, stream=stream) + + @staticmethod + def current_stream(device: Any) -> int: + import torch + + return int(torch.cuda.current_stream(device=device).cuda_stream) + + def _register_one(self, ptr: int, length: int, generation: int, *, read_only: bool) -> None: + if ptr <= 0: + raise ValueError(f"Fluxon segment pointer must be positive, got {ptr}") + key = (ptr, length, generation) + with self._register_lock: + if key in self._registered: + return + base_flags = int(self._api.cudaHostRegisterPortable) | int( + self._api.cudaHostRegisterMapped + ) + flags = base_flags + if read_only: + flags |= int(self._api.cudaHostRegisterReadOnly) + code = self._code(self._api.cudaHostRegister(ptr, length, flags)) + success = int(self._enum("cudaError_t", "cudaSuccess")) + already = int( + self._enum("cudaError_t", "cudaErrorHostMemoryAlreadyRegistered") + ) + if code not in (success, already) and read_only: + code = self._code(self._api.cudaHostRegister(ptr, length, base_flags)) + if code not in (success, already): + raise RuntimeError( + f"cudaHostRegister failed with CUDA error {code}: " + f"ptr={ptr:#x} len={length} generation={generation}" + ) + self._registered.add(key) + + def register_fluxon_segments(self, segments: Iterable[Mapping[str, Any]]) -> None: + for segment in segments: + length = int(segment["len"]) + generation = int(segment["generation"]) + if length <= 0: + raise ValueError(f"Fluxon segment len must be positive, got {length}") + self._register_one(int(segment["write_ptr"]), length, generation, read_only=False) + if "read_ptr" in segment: + self._register_one(int(segment["read_ptr"]), length, generation, read_only=True) diff --git a/cacheseek/stores/fluxon.py b/cacheseek/stores/fluxon.py index 05ad4e1..276f302 100644 --- a/cacheseek/stores/fluxon.py +++ b/cacheseek/stores/fluxon.py @@ -7,17 +7,18 @@ ## External dependencies -- The ``fluxon_py`` package (not pip-installable; provided by the TeleAI Fluxon - repo as pylib_src source or a prebuilt wheel). +- The ``fluxon_py`` package, provided by a pinned TeleAI Fluxon wheel/package or + an explicitly configured editable/source installation. - A Fluxon YAML config (``fluxon_config_path``) is required before startup, specifying instance_key, shared_memory_path, cluster_name, etc. See the deployment examples in the Fluxon repo. ## Import strategy (lazy) -1. Prefer local source under ``fluxon/pylib_src/`` at the repo root. -2. Fall back to the installed ``fluxon_py`` package. -3. If both fail, raise ``ImportError`` at the call site. +1. Resolve ``fluxon_py`` through Python's normal package import rules. +2. Production deployments should pin the Fluxon wheel/package. Source-based + development should use an editable install or an explicit ``PYTHONPATH``. +3. Log the resolved Python package and loaded native-extension identities. The module itself always imports as long as ``__init__`` is not actually invoked (see ``__init__.py``). @@ -51,15 +52,43 @@ from __future__ import annotations +import ctypes +import hashlib +import importlib +import importlib.machinery +import sys +from functools import cache +from pathlib import Path from typing import Any from loguru import logger +@cache +def _file_sha256(path: str) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _resolved_module_path(module: Any) -> Path | None: + raw_path = getattr(module, "__file__", None) + if not raw_path: + return None + path = Path(raw_path) + try: + return path.resolve() + except OSError: + return path.absolute() + + class FluxonKVStore: """Fluxon KV store adapter. See the module docstring for the error policy.""" _BYTES_FIELD_KEY = "v" # Fluxon examples/tests use "v" as the bytes field key. + _PLAN_BLOB_MAGIC = 0x4658_504C_414E_5631 def __init__( self, @@ -86,13 +115,17 @@ def __init__( self._store = store return if not config_path: - raise ValueError( - "FluxonKVStore requires either config_path or an injected store." - ) + raise ValueError("FluxonKVStore requires either config_path or an injected store.") fluxon = self._import_fluxon_layer() cfg = fluxon.FluxonKvClientConfig.from_file(config_path) - res = fluxon.new_store(cfg) + try: + res = fluxon.new_store(cfg) + finally: + # new_store() is the point at which Fluxon loads fluxon_pyo3. Inspect + # only modules that Fluxon actually loaded; importing the extension + # ourselves would bypass its loader and venv-authority checks. + self._log_fluxon_native_identity() store_obj = self._unwrap_result_ok(res, op="new_store") if store_obj is None: err = self._unwrap_result_err(res, op="new_store") @@ -180,120 +213,131 @@ def remove(self, key: str) -> None: def list_keys(self) -> list[str]: """Fluxon offers no key-enumeration interface; returns an empty list.""" logger.warning( - "FluxonKVStore.list_keys() is not supported by the Fluxon backend; " - "returning []." + "FluxonKVStore.list_keys() is not supported by the Fluxon backend; returning []." ) return [] # ------------------------------------------------------------------ - # TensorKVStore capability — DLPack zero-pickle tensor put/get. + # fluxon_new plan APIs. # - # put_tensor hands the tensor's DLPack pointer to Fluxon (no Python bytes, no - # pickle); on get_tensor, access() returns a DLPack view onto the mapped shared - # memory, and from_dlpack + clone materializes an owned CPU tensor (decoupled - # from the MemHolder lifetime). DLPack supports only CPU, C-contiguous tensors, - # hence .detach().cpu().contiguous() before put. + # get_start/get_transfer expose read-side prefix plans. local_fast_put_start + # exposes write-side local writable plans. The returned plan_ptr points to a + # u64 blob: [magic, value_ptr_count, value_ptr...]. The raw Fluxon store owns + # the registry entry behind the plan; CacheSeek must release/commit/abort it. # ------------------------------------------------------------------ - def put_tensor(self, key: str, tensor: Any) -> None: - """Store ``tensor`` under ``key`` via DLPack — no pickle, no Python bytes. - - The tensor is moved to CPU and made C-contiguous (DLPack constraints) - before its DLPack pointer is handed to Fluxon. Any backend failure is - wrapped and re-raised as RuntimeError. - """ - import torch # noqa: F401 (lazy — heavy dep stays out of import time) - + def wait_local_segments_ready(self) -> list[dict[str, Any]]: + """Return the local owner mappings used by direct CUDA transfers.""" try: - t = tensor.detach().to("cpu").contiguous() - res = self._store.put(key, {self._BYTES_FIELD_KEY: t}) - fut = self._unwrap_result_ok(res, op="put_tensor") - if fut is None: - err = self._unwrap_result_err(res, op="put_tensor") - raise RuntimeError(err) - wait_res = fut.wait() - if self._unwrap_result_ok(wait_res, op="put_tensor.wait") is None: - err = self._unwrap_result_err(wait_res, op="put_tensor.wait") - raise RuntimeError(err) + segments = self._store.wait_local_segments_ready() + if not isinstance(segments, list) or not segments: + raise RuntimeError("Fluxon returned no local segments") + return [dict(segment) for segment in segments] except Exception as exc: - logger.exception("FluxonKV.put_tensor failed key={} err={}", key, exc) - raise RuntimeError( - f"FluxonKV.put_tensor failed key={key!r}: {exc}" - ) from exc + raise RuntimeError(f"FluxonKV.wait_local_segments_ready failed: {exc}") from exc - def get_tensor( + def get_start( self, - key: str, + keys: list[str], *, - shape: tuple[int, ...], - dtype: Any, - device: Any = None, - ) -> Any | None: - """Fetch a tensor for ``key``, or None on a cache miss. - - On a DLPack payload, materializes an owned CPU tensor via - ``from_dlpack(...).clone()`` so it outlives the backend MemHolder. On a - raw-bytes payload, reinterprets the bytes using ``dtype`` and ``shape``. - ``device``, if given, moves the result onto that device. + prefix_best_effort: bool = True, + atomic_group_lens: list[int] | None = None, + ) -> Any: + try: + return self._store.get_start( + keys, + prefix_best_effort=prefix_best_effort, + atomic_group_lens=atomic_group_lens, + ) + except Exception as exc: + logger.exception("FluxonKV.get_start failed keys={} err={}", len(keys), exc) + raise RuntimeError(f"FluxonKV.get_start failed keys={len(keys)}: {exc}") from exc - Args: - key: Store key to read. - shape: Target tensor shape, used to reshape a raw-bytes payload. - dtype: torch dtype used to reinterpret a raw-bytes payload. - device: Optional device to move the returned tensor onto. + def get_transfer(self, handle: Any) -> int: + try: + return int(self._store.get_transfer(handle)) + except Exception as exc: + logger.exception("FluxonKV.get_transfer failed err={}", exc) + raise RuntimeError(f"FluxonKV.get_transfer failed: {exc}") from exc - Raises: - RuntimeError: Any backend failure other than a key-not-found miss - (which returns None). Every error — including an unexpected - payload type — is wrapped as RuntimeError, never propagated raw. - """ - import torch + def cancel_get_transfer(self, handle: Any) -> None: + try: + self._store.cancel_get_transfer(handle) + except Exception as exc: + logger.exception("FluxonKV.cancel_get_transfer failed err={}", exc) + raise RuntimeError(f"FluxonKV.cancel_get_transfer failed: {exc}") from exc + def release_views(self, plan_ptr: int) -> None: try: - r = self._store.get(key) - fut = self._unwrap_result_ok(r, op="get_tensor") - if fut is None: - err = self._unwrap_result_err(r, op="get_tensor") - if self._is_key_not_found(err): - return None - raise RuntimeError(err) + self._store.release_views(int(plan_ptr)) + except Exception as exc: + logger.exception("FluxonKV.release_views failed plan_ptr={} err={}", plan_ptr, exc) + raise RuntimeError(f"FluxonKV.release_views failed plan_ptr={plan_ptr}: {exc}") from exc - wait_res = fut.wait() - mh = self._unwrap_result_ok(wait_res, op="get_tensor.wait") - if mh is None: - err = self._unwrap_result_err(wait_res, op="get_tensor.wait") - if self._is_key_not_found(err): - return None - raise RuntimeError(err) + def local_fast_put_start( + self, + keys: list[str], + value_len: int, + opts: Any | None = None, + ) -> int: + try: + return int(self._store.local_fast_put_start(keys, int(value_len), opts)) + except Exception as exc: + logger.exception( + "FluxonKV.local_fast_put_start failed keys={} value_len={} err={}", + len(keys), + value_len, + exc, + ) + raise RuntimeError( + f"FluxonKV.local_fast_put_start failed keys={len(keys)} " + f"value_len={value_len}: {exc}" + ) from exc - d_res = mh.access() - if not d_res.is_ok(): - err = d_res.unwrap_error("get_tensor.access failed") - raise RuntimeError(f"err_type={type(err).__name__} err={err}") - d = d_res.unwrap("get_tensor.access ok") - v = d.get(self._BYTES_FIELD_KEY) - if hasattr(v, "__dlpack__"): - # zero-copy view into the pool; clone to own it past the holder. - t = torch.from_dlpack(v).clone() - elif isinstance(v, (bytes, bytearray, memoryview)): - t = ( - torch.frombuffer(bytes(v), dtype=torch.uint8) - .view(dtype) - .reshape(shape) - ) - else: - raise TypeError( - f"Fluxon get_tensor payload field {self._BYTES_FIELD_KEY!r} is " - f"neither dlpack nor bytes (type={type(v).__name__})" - ) - if device is not None: - t = t.to(device) - return t + def local_fast_put_commit(self, plan_ptr: int) -> Any: + try: + return self._store.local_fast_put_commit(int(plan_ptr)) except Exception as exc: - logger.exception("FluxonKV.get_tensor failed key={} err={}", key, exc) + logger.exception( + "FluxonKV.local_fast_put_commit failed plan_ptr={} err={}", plan_ptr, exc + ) raise RuntimeError( - f"FluxonKV.get_tensor failed key={key!r}: {exc}" + f"FluxonKV.local_fast_put_commit failed plan_ptr={plan_ptr}: {exc}" ) from exc + def put_abort(self, plan_ptr: int) -> None: + try: + self._store.put_abort(int(plan_ptr)) + except Exception as exc: + logger.exception("FluxonKV.put_abort failed plan_ptr={} err={}", plan_ptr, exc) + raise RuntimeError(f"FluxonKV.put_abort failed plan_ptr={plan_ptr}: {exc}") from exc + + @classmethod + def decode_plan_ptr(cls, plan_ptr: int, expected_count: int) -> list[int]: + """Decode a fluxon_new plan blob into value pointers. + + The Fluxon pyo3 layer builds the blob as ``u64[magic, count, ptr...]``. + ``plan_ptr`` is the address of the first u64. The caller owns no memory; + it must keep the Fluxon plan alive until all returned pointers are done. + """ + if plan_ptr <= 0: + raise ValueError(f"plan_ptr must be positive, got {plan_ptr!r}") + if expected_count < 0: + raise ValueError(f"expected_count must be >= 0, got {expected_count!r}") + header = (ctypes.c_uint64 * 2).from_address(int(plan_ptr)) + magic = int(header[0]) + count = int(header[1]) + if magic != cls._PLAN_BLOB_MAGIC: + raise ValueError( + f"Fluxon plan blob magic mismatch: expected=0x{cls._PLAN_BLOB_MAGIC:x} " + f"got=0x{magic:x}" + ) + if count != expected_count: + raise ValueError( + f"Fluxon plan pointer count mismatch: expected={expected_count} got={count}" + ) + words = (ctypes.c_uint64 * (2 + count)).from_address(int(plan_ptr)) + return [int(words[2 + i]) for i in range(count)] + def _extract_bytes_from_value( self, value_obj: Any, @@ -319,7 +363,7 @@ def _extract_bytes_from_value( """ if value_obj is None: return None - if isinstance(value_obj, (bytes, bytearray, memoryview)): + if isinstance(value_obj, bytes | bytearray | memoryview): return bytes(value_obj) context = f" key={key!r}" if key is not None else "" @@ -333,9 +377,7 @@ def _extract_bytes_from_value( d_res = access() except Exception as exc: logger.exception("Fluxon MemHolder.access failed{} err={}", context, exc) - raise RuntimeError( - f"Fluxon MemHolder.access failed{context}: {exc}" - ) from exc + raise RuntimeError(f"Fluxon MemHolder.access failed{context}: {exc}") from exc if not d_res.is_ok(): err = d_res.unwrap_error("mem.access failed") @@ -352,7 +394,7 @@ def _extract_bytes_from_value( ) v = d.get(self._BYTES_FIELD_KEY) - if not isinstance(v, (bytes, bytearray, memoryview)): + if not isinstance(v, bytes | bytearray | memoryview): raise TypeError( "Fluxon MemHolder.access payload does not contain bytes field " f"field={self._BYTES_FIELD_KEY!r} value_type={type(v).__name__}" @@ -392,48 +434,71 @@ def _is_key_not_found(self, err: Any) -> bool: def _import_fluxon_layer(self): """Lazily import the Fluxon Python API. - Prefer local source under ``fluxon/pylib_src/`` at the repo root (easier to - debug), falling back to the installed ``fluxon_py`` package. Raise - ImportError if both fail. + Package selection is delegated to the active Python environment. CacheSeek + deliberately does not scan workspace directories or modify ``sys.path``. """ - load_errors: list[str] = [] - - # 1) Prefer local Fluxon source if present. try: - import sys - from pathlib import Path as _Path - - # __file__ -> stores/fluxon.py; parents[3] -> repo root - repo_root = _Path(__file__).resolve().parents[3] - fluxon_root = ( - repo_root / "fluxon" / "pylib_src" - ) # the fluxon_py package lives under pylib_src - if fluxon_root.is_dir(): - fluxon_path = str(fluxon_root) - if fluxon_path not in sys.path: - sys.path.insert(0, fluxon_path) - import fluxon_py as api # type: ignore - - return api + api = importlib.import_module("fluxon_py") except Exception as exc: - load_errors.append( - f"local Fluxon source import failed type={type(exc).__name__} err={exc}" - ) - - # 2) Fall back to the installed fluxon_py package. - try: - import fluxon_py as api # type: ignore + raise ImportError( + "Fluxon Python API import failed (fluxon_py) " + f"type={type(exc).__name__} err={exc}. Install a pinned Fluxon " + "wheel/package; for source development, use an editable install " + "or put the selected source root on PYTHONPATH explicitly." + ) from exc - return api - except Exception as exc: - load_errors.append( - f"installed fluxon_py import failed type={type(exc).__name__} err={exc}" + module_path = _resolved_module_path(api) + logger.info( + "Fluxon Python package identity module={} version={} path={}", + getattr(api, "__name__", "fluxon_py"), + getattr(api, "__version__", ""), + str(module_path) if module_path is not None else "", + ) + return api + + @staticmethod + def _loaded_fluxon_native_extensions() -> list[tuple[str, Path]]: + suffixes = tuple(importlib.machinery.EXTENSION_SUFFIXES) + identities: list[tuple[str, Path]] = [] + seen_paths: set[Path] = set() + for module_name, module in sorted(list(sys.modules.items())): + if module_name != "fluxon_pyo3" and not module_name.startswith("fluxon_pyo3."): + continue + module_path = _resolved_module_path(module) + if module_path is None or not str(module_path).endswith(suffixes): + continue + if module_path in seen_paths: + continue + seen_paths.add(module_path) + identities.append((module_name, module_path)) + return identities + + @classmethod + def _log_fluxon_native_identity(cls) -> None: + identities = cls._loaded_fluxon_native_extensions() + if not identities: + logger.warning( + "Fluxon native extension identity unavailable: " + "no loaded fluxon_pyo3 extension module" ) + return - detail = ( - "; ".join(load_errors) if load_errors else "no import attempts succeeded" - ) - raise ImportError( - f"Fluxon Python API not available (fluxon_py): {detail}. " - "Install fluxon_py or place the source at /fluxon/pylib_src/." - ) + for module_name, module_path in identities: + try: + sha256 = _file_sha256(str(module_path)) + except OSError as exc: + logger.warning( + "Fluxon native extension identity module={} path={} " + "sha256= err_type={} err={}", + module_name, + str(module_path), + type(exc).__name__, + exc, + ) + continue + logger.info( + "Fluxon native extension identity module={} path={} sha256={}", + module_name, + str(module_path), + sha256, + ) diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index 2d85263..e33d3e5 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -5,9 +5,9 @@ - InMemoryTierStore: in-process dict. ``put_async`` runs synchronously and calls ``on_ready`` immediately. For end-to-end integration and tests; no real data movement, no tiering, no async queue. -- TensorStoreTierStore: adapts any duck-typed tensor store (``put_tensor(key, t)`` - / ``get_tensor(key) -> t``, e.g. Fluxon TensorKVStore) into a KVTierStore, - storing each layer's payload under its own key. +- TensorStoreTierStore: adapts a generic per-key tensor backend into a KVTierStore. +- PlanTensorTierStore: synchronously transfers raw GPU tensor bytes through + Fluxon Plan pointers after registering the owner mappings once. """ from __future__ import annotations @@ -17,12 +17,21 @@ import queue import threading from collections.abc import Callable, Sequence +from dataclasses import dataclass from pathlib import Path from typing import Any from loguru import logger from .base import BlobHandle, Tier +from .cuda_transfer import CudaTransferRuntime, PointerCopy + + +@dataclass(frozen=True, slots=True) +class _GpuPlanEntry: + key: str + tensor: Any + nbytes: int class InMemoryTierStore: @@ -67,9 +76,8 @@ class LocalDiskTensorStore: with TensorStoreTierStore). Raw bytes, one file per key (no pickle overhead); the bytes ``put/get`` path - serves the ``:spec`` sidecar. It goes through the same adapter as the Fluxon - backend (async writes, spec bookkeeping), so the backend is the only variable - when comparing against a baseline. + serves the ``:spec`` sidecar. It uses the generic asynchronous tensor adapter; + Fluxon exact-prefix traffic uses the separate synchronous Plan adapter. Note: writes go through the page cache (no fsync), matching ordinary file-server semantics; sustained throughput is still disk-bound. """ @@ -112,11 +120,9 @@ def get(self, key: str) -> bytes | None: class TensorStoreTierStore: - """Adapts any ``put_tensor/get_tensor``-style tensor store (e.g. Fluxon - TensorKVStore). + """Adapts a generic ``put_tensor/get_tensor``-style tensor store. - Fluxon's ``get_tensor(key, *, shape, dtype)`` requires those keyword-only args: - it stores a raw buffer, so reads must supply the view spec. Therefore each + Raw tensor buffers need shape and dtype when reconstructed. Therefore each key's (shape, dtype) is recorded on put — in an in-process dict as the primary path, and, if the backend exposes bytes ``put/get``, also written to a ``:spec`` sidecar (JSON) so the spec survives across processes. @@ -172,7 +178,7 @@ def _get_one(self, key: str) -> Any: # ----------------------------------------------------------------- async write def _do_put_payload(self, locator: str, payload: Sequence[Any]) -> None: for i, p in enumerate(payload): - if isinstance(p, (tuple, list)) and len(p) == 2: + if isinstance(p, tuple | list) and len(p) == 2: # per-layer payload = (k, v) tensor pair -> two keys (put_tensor # takes a single tensor) self._put_one(self._layer_key(locator, i) + ":k", p[0]) @@ -240,3 +246,284 @@ def get_skeleton(self, locator: str) -> Any: def free(self, handle: BlobHandle) -> None: # noqa: ARG002 -- no-op (reclamation not implemented) """No-op: reclamation against the backing tensor store is not implemented.""" return None + + +class PlanTensorTierStore: + """Synchronous Fluxon Plan adapter for exact-prefix GPU tensor bytes. + + Writes use local writable value pointers; reads keep source views alive until + the single H2D batch completes. Owner mappings are registered once per store. + """ + + def __init__(self, plan_store: Any, *, cuda_runtime: Any | None = None) -> None: + self._ps = plan_store + self._cuda = CudaTransferRuntime() if cuda_runtime is None else cuda_runtime + self._segments_ready = False + self._segments_lock = threading.Lock() + + @staticmethod + def _layer_key(locator: str, layer: int) -> str: + return f"{locator}:L{layer}" + + @staticmethod + def _tensor_nbytes(tensor: Any) -> int: + return int(tensor.numel()) * int(tensor.element_size()) + + def _ensure_segments_ready(self) -> None: + if self._segments_ready: + return + with self._segments_lock: + if self._segments_ready: + return + segments = self._ps.wait_local_segments_ready() + self._cuda.register_fluxon_segments(segments) + self._segments_ready = True + + def _flatten_payload(self, locator: str, payload: Sequence[Any]) -> list[_GpuPlanEntry]: + entries: list[_GpuPlanEntry] = [] + for i, p in enumerate(payload): + if isinstance(p, tuple | list) and len(p) == 2: + for suffix, tensor in ((":k", p[0]), (":v", p[1])): + key = self._layer_key(locator, i) + suffix + entries.append( + _GpuPlanEntry(key=key, tensor=tensor, nbytes=self._tensor_nbytes(tensor)) + ) + else: + key = self._layer_key(locator, i) + entries.append(_GpuPlanEntry(key=key, tensor=p, nbytes=self._tensor_nbytes(p))) + return entries + + @staticmethod + def _group_by_nbytes(entries: Sequence[_GpuPlanEntry]) -> dict[int, list[_GpuPlanEntry]]: + by_nbytes: dict[int, list[_GpuPlanEntry]] = {} + for entry in entries: + by_nbytes.setdefault(entry.nbytes, []).append(entry) + return by_nbytes + + @staticmethod + def _validate_gpu_entries(entries: Sequence[_GpuPlanEntry]) -> Any: + if not entries: + raise ValueError("Fluxon Plan put requires at least one tensor") + device = entries[0].tensor.device + for entry in entries: + tensor = entry.tensor + if not tensor.is_cuda: + raise ValueError(f"Fluxon Plan put requires CUDA tensors: key={entry.key!r}") + if not tensor.is_contiguous(): + raise ValueError(f"Fluxon Plan put requires contiguous tensors: key={entry.key!r}") + if tensor.device != device: + raise ValueError( + f"Fluxon Plan put requires one CUDA device: expected={device} got={tensor.device}" + ) + actual_nbytes = PlanTensorTierStore._tensor_nbytes(tensor) + if actual_nbytes != entry.nbytes or actual_nbytes <= 0: + raise ValueError( + f"invalid tensor byte size for key={entry.key!r}: " + f"expected={entry.nbytes} got={actual_nbytes}" + ) + return device + + @staticmethod + def _wait_commit(future: Any) -> None: + wait_result = future.wait() + ok = wait_result.is_ok() if hasattr(wait_result, "is_ok") else bool(wait_result) + if not ok: + error = ( + wait_result.unwrap_error("local_fast_put_commit failed") + if hasattr(wait_result, "unwrap_error") + else wait_result + ) + raise RuntimeError(f"local_fast_put_commit failed: {error}") + if hasattr(wait_result, "unwrap"): + ret_codes = wait_result.unwrap("local_fast_put_commit ok") + if isinstance(ret_codes, list): + failed = [int(code) for code in ret_codes if int(code) != 0] + if failed: + raise RuntimeError( + f"local_fast_put_commit returned nonzero ret codes: {failed}" + ) + + def _put_entries(self, entries: Sequence[_GpuPlanEntry]) -> None: + device = self._validate_gpu_entries(entries) + self._ensure_segments_ready() + stream = self._cuda.current_stream(device) + + plans: list[int] = [] + submitted: set[int] = set() + copies: list[PointerCopy] = [] + try: + for value_len, group in self._group_by_nbytes(entries).items(): + plan_ptr = self._ps.local_fast_put_start( + [entry.key for entry in group], value_len + ) + plans.append(plan_ptr) + destinations = self._ps.decode_plan_ptr(plan_ptr, expected_count=len(group)) + copies.extend( + PointerCopy( + dst=destination, + src=int(entry.tensor.data_ptr()), + nbytes=value_len, + ) + for destination, entry in zip(destinations, group, strict=True) + ) + self._cuda.copy_d2h(copies, stream=stream) + for plan_ptr in plans: + future = self._ps.local_fast_put_commit(plan_ptr) + submitted.add(plan_ptr) + self._wait_commit(future) + except Exception: + for plan_ptr in plans: + if plan_ptr not in submitted: + with contextlib.suppress(Exception): + self._ps.put_abort(plan_ptr) + raise + + def flush(self) -> None: + return None + + def put_async( + self, + locator: str, + payload: Sequence[Any], + *, + tier: Tier, + on_ready: Callable[[], None] | None = None, + ) -> None: + del tier + self._put_entries(self._flatten_payload(locator, payload)) + if on_ready is not None: + on_ready() + + def put_skeleton(self, locator: str, latent: Any) -> None: + self._put_entries( + [_GpuPlanEntry(key=locator, tensor=latent, nbytes=self._tensor_nbytes(latent))] + ) + + def materialize_skeletons(self, targets: Sequence[tuple[str, Any]]) -> bool: + """Restore all skeleton latents in one Plan GET and one H2D batch.""" + if not targets: + return False + entries = [ + _GpuPlanEntry(key=locator, tensor=target, nbytes=self._tensor_nbytes(target)) + for locator, target in targets + ] + device = self._validate_gpu_entries(entries) + self._ensure_segments_ready() + stream = self._cuda.current_stream(device) + keys = [entry.key for entry in entries] + transfer_handle = self._ps.get_start( + keys, + prefix_best_effort=False, + atomic_group_lens=[1] * len(keys), + ) + plan_ptr: int | None = None + try: + if transfer_handle.result.transferable_len != len(keys): + self._ps.cancel_get_transfer(transfer_handle) + return False + plan_ptr = self._ps.get_transfer(transfer_handle) + sources = self._ps.decode_plan_ptr(plan_ptr, expected_count=len(keys)) + copies = [ + PointerCopy( + dst=int(entry.tensor.data_ptr()), + src=source, + nbytes=entry.nbytes, + ) + for source, entry in zip(sources, entries, strict=True) + ] + self._cuda.copy_h2d(copies, stream=stream) + return True + except Exception: + if plan_ptr is None and not getattr(transfer_handle, "closed", False): + with contextlib.suppress(Exception): + self._ps.cancel_get_transfer(transfer_handle) + raise + finally: + if plan_ptr is not None: + self._ps.release_views(plan_ptr) + + def materialize_path(self, path: Sequence[Any], window: Any, *, depth: int) -> bool: + if not path: + return False + last_blob = path[-1].blob + if last_blob is None or not last_blob.ready: + return False + n_layers = int(last_blob.n_layers) + if n_layers <= 0: + raise ValueError(f"invalid Fluxon Plan layer count: {n_layers}") + + expected_chunk_nbytes = sum( + 2 * int(window.chunk_value_nbytes(layer)) for layer in range(n_layers) + ) + keys: list[str] = [] + for node in path: + blob = node.blob + if blob is None or not blob.ready: + return False + if int(blob.n_layers) != n_layers or int(blob.nbytes) != expected_chunk_nbytes: + raise ValueError( + "Fluxon Plan KV layout mismatch: " + f"locator={blob.locator!r} expected_layers={n_layers} " + f"actual_layers={blob.n_layers} expected_nbytes={expected_chunk_nbytes} " + f"actual_nbytes={blob.nbytes}" + ) + for layer in range(n_layers): + keys.extend( + [ + self._layer_key(blob.locator, layer) + ":k", + self._layer_key(blob.locator, layer) + ":v", + ] + ) + + self._ensure_segments_ready() + transfer_handle = self._ps.get_start( + keys, + prefix_best_effort=True, + atomic_group_lens=[n_layers * 2] * len(path), + ) + plan_ptr: int | None = None + try: + if transfer_handle.result.transferable_len != len(keys): + self._ps.cancel_get_transfer(transfer_handle) + return False + plan_ptr = self._ps.get_transfer(transfer_handle) + decoded = self._ps.decode_plan_ptr(plan_ptr, expected_count=len(keys)) + per_layer: list[list[tuple[int, tuple[int, int]]]] = [ + [] for _ in range(n_layers) + ] + cursor = 0 + for node in path: + for layer in range(n_layers): + k_ptr, v_ptr = decoded[cursor], decoded[cursor + 1] + cursor += 2 + per_layer[layer].append((node.depth, (k_ptr, v_ptr))) + if cursor != len(decoded): + raise RuntimeError( + f"unused Fluxon pointers: consumed={cursor} total={len(decoded)}" + ) + + copies: list[PointerCopy] = [] + for layer, blobs in enumerate(per_layer): + copies.extend(window.build_seed_copies(layer, blobs, depth)) + stream = self._cuda.current_stream(window.cuda_device) + self._cuda.copy_h2d(copies, stream=stream) + window.set_resume_depth(depth) + return True + except Exception: + if plan_ptr is None and not getattr(transfer_handle, "closed", False): + with contextlib.suppress(Exception): + self._ps.cancel_get_transfer(transfer_handle) + raise + finally: + if plan_ptr is not None: + self._ps.release_views(plan_ptr) + + def free(self, handle: BlobHandle) -> None: + """Remove a blob's layer K/V values while retaining its skeleton latent.""" + n_layers = int(handle.n_layers) + if n_layers <= 0: + raise ValueError(f"invalid Fluxon Plan layer count: {n_layers}") + for layer in range(n_layers): + layer_key = self._layer_key(handle.locator, layer) + self._ps.remove(layer_key + ":k") + self._ps.remove(layer_key + ":v") diff --git a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py index 0a4f2e2..99ad1fd 100644 --- a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py +++ b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py @@ -161,9 +161,8 @@ def run_request(pipeline, binding, frame_num: int, poses, intrinsics, *, image_p frames.extend(pipeline.generate_next_chunk(runtime)) chunk_times.append(round(time.time() - tc, 3)) wall = round(time.time() - t0, 3) - # Async writes: drain in-flight puts so the next request's lookup sees all of this - # request's chunks. flush_s = write time not hidden by the chunk loop; smaller means - # more of the async benefit was realized. + # Give each store a request-boundary visibility point. The synchronous Fluxon + # Plan adapter treats flush() as a no-op; the LocalDisk worker drains its queue. tf = time.time() store = binding.mgr.store if hasattr(store, "flush"): @@ -224,17 +223,15 @@ def main() -> int: def make_store(): if args.store == "fluxon": - from cacheseek.stores import TensorStoreTierStore + from cacheseek.stores import PlanTensorTierStore from cacheseek.stores.fluxon import FluxonKVStore assert args.fluxon_config, "--store fluxon requires --fluxon-config" - # True async writes: put enqueues and returns, worker publishes ready after - # draining (write latency moved out of the chunk loop). - return TensorStoreTierStore(FluxonKVStore(config_path=args.fluxon_config), async_put=True) + # Fluxon Plan writes synchronously submit one GPU-to-segment batch. + return PlanTensorTierStore(FluxonKVStore(config_path=args.fluxon_config)) if args.store == "localdisk": from cacheseek.stores import TensorStoreTierStore from cacheseek.stores.tier import LocalDiskTensorStore - # Same adapter and same async path as fluxon, so the only variable in the - # baseline comparison is the backend. + # LocalDisk retains the generic asynchronous tensor adapter. return TensorStoreTierStore(LocalDiskTensorStore(args.disk_root), async_put=True) return InMemoryTierStore() diff --git a/pyproject.toml b/pyproject.toml index 50e22e7..19f48a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ qdrant = ["qdrant-client>=1.7"] # vector store backend faiss = ["faiss-cpu>=1.7"] # vector store backend encoder = ["transformers>=4.57", "accelerate>=0.30", "qwen-vl-utils", "scipy", "sentencepiece"] # Qwen3-VL embed/rerank backend (qwen_vl_utils->vision_process, scipy->reranker softmax, sentencepiece->Qwen tokenizer) +fluxon = ["cuda-bindings>=12.6"] dev = ["pytest>=7", "pytest-asyncio>=0.21", "ruff", "pyright", "httpx>=0.27"] # Everything needed to run the full reference path (all backends). all = ["cacheseek[qdrant,faiss,encoder]"] @@ -40,6 +41,8 @@ asyncio_mode = "auto" markers = [ "smoke: fast import/config/script checks that validate local package wiring", "e2e: end-to-end lifecycle checks; external variants may require opt-in infrastructure", + "cuda: requires a real CUDA device and cuda-bindings", + "fluxon: requires a real Fluxon pyo3 store and config", ] [tool.ruff] diff --git a/tests/test_cuda_transfer.py b/tests/test_cuda_transfer.py new file mode 100644 index 0000000..ab143b1 --- /dev/null +++ b/tests/test_cuda_transfer.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import pytest + +from cacheseek.stores.cuda_transfer import CudaTransferRuntime, PointerCopy + + +class FakeCudaBinding: + cudaSuccess = 0 + cudaErrorHostMemoryAlreadyRegistered = 712 + cudaMemcpyHostToDevice = 1 + cudaMemcpyDeviceToHost = 2 + cudaHostRegisterPortable = 1 + cudaHostRegisterMapped = 2 + cudaHostRegisterReadOnly = 8 + + def __init__(self) -> None: + self.register_calls: list[tuple[int, int, int]] = [] + self.copy_calls: list[tuple[int, int, int, int, int]] = [] + self.sync_calls: list[int] = [] + self.register_results: list[int | tuple[int]] = [] + self.copy_results: list[int | tuple[int]] = [] + self.sync_results: list[int | tuple[int]] = [] + + def cudaHostRegister(self, ptr: int, size: int, flags: int) -> int | tuple[int]: + self.register_calls.append((int(ptr), int(size), int(flags))) + if self.register_results: + return self.register_results.pop(0) + return self.cudaSuccess + + def cudaMemcpyAsync( + self, dst: int, src: int, nbytes: int, kind: int, stream: int + ) -> int | tuple[int]: + self.copy_calls.append((int(dst), int(src), int(nbytes), int(kind), int(stream))) + if self.copy_results: + return self.copy_results.pop(0) + return (self.cudaSuccess,) + + def cudaStreamSynchronize(self, stream: int) -> int | tuple[int]: + self.sync_calls.append(int(stream)) + if self.sync_results: + return self.sync_results.pop(0) + return self.cudaSuccess + + +def test_copy_d2h_submits_all_descriptors_then_syncs_once() -> None: + api = FakeCudaBinding() + cuda = CudaTransferRuntime(api=api) + + cuda.copy_d2h( + [PointerCopy(dst=100, src=1000, nbytes=16), PointerCopy(dst=200, src=2000, nbytes=32)], + stream=77, + ) + + assert api.copy_calls == [(100, 1000, 16, 2, 77), (200, 2000, 32, 2, 77)] + assert api.sync_calls == [77] + + +def test_copy_h2d_uses_host_to_device_direction() -> None: + api = FakeCudaBinding() + cuda = CudaTransferRuntime(api=api) + + cuda.copy_h2d([PointerCopy(dst=1000, src=100, nbytes=16)], stream=88) + + assert api.copy_calls == [(1000, 100, 16, 1, 88)] + assert api.sync_calls == [88] + + +def test_registers_write_and_read_mapping_once() -> None: + api = FakeCudaBinding() + cuda = CudaTransferRuntime(api=api) + segments = [{"write_ptr": 10_001, "read_ptr": 20_001, "len": 4096, "generation": 9}] + + cuda.register_fluxon_segments(segments) + cuda.register_fluxon_segments(segments) + + assert [call[:2] for call in api.register_calls] == [(10_001, 4096), (20_001, 4096)] + + +def test_already_registered_is_success() -> None: + api = FakeCudaBinding() + api.register_results = [(api.cudaErrorHostMemoryAlreadyRegistered,)] * 2 + cuda = CudaTransferRuntime(api=api) + + cuda.register_fluxon_segments( + [{"write_ptr": 10_002, "read_ptr": 20_002, "len": 1024, "generation": 10}] + ) + + assert len(api.register_calls) == 2 + + +def test_read_only_registration_retries_without_read_only() -> None: + api = FakeCudaBinding() + api.register_results = [api.cudaSuccess, 999, (api.cudaSuccess,)] + cuda = CudaTransferRuntime(api=api) + + cuda.register_fluxon_segments( + [{"write_ptr": 10_003, "read_ptr": 20_003, "len": 2048, "generation": 11}] + ) + + base_flags = api.cudaHostRegisterPortable | api.cudaHostRegisterMapped + assert api.register_calls == [ + (10_003, 2048, base_flags), + (20_003, 2048, base_flags | api.cudaHostRegisterReadOnly), + (20_003, 2048, base_flags), + ] + + +def test_memcpy_failure_syncs_previously_submitted_work() -> None: + api = FakeCudaBinding() + api.copy_results = [api.cudaSuccess, (700,)] + cuda = CudaTransferRuntime(api=api) + + with pytest.raises(RuntimeError, match=r"cudaMemcpyAsync\[1\].*700"): + cuda.copy_d2h( + [PointerCopy(dst=100, src=1000, nbytes=16), PointerCopy(dst=200, src=2000, nbytes=32)], + stream=77, + ) + + assert api.sync_calls == [77] + + +def test_invalid_descriptor_is_rejected_before_any_memcpy() -> None: + api = FakeCudaBinding() + cuda = CudaTransferRuntime(api=api) + + with pytest.raises(ValueError, match=r"invalid pointer copy at index 1"): + cuda.copy_d2h( + [PointerCopy(dst=100, src=1000, nbytes=16), PointerCopy(dst=0, src=2000, nbytes=32)], + stream=77, + ) + + assert api.copy_calls == [] + assert api.sync_calls == [] + + +def test_memcpy_and_cleanup_sync_failures_are_both_reported() -> None: + api = FakeCudaBinding() + api.copy_results = [api.cudaSuccess, (700,)] + api.sync_results = [(701,)] + cuda = CudaTransferRuntime(api=api) + + with pytest.raises( + RuntimeError, + match=r"cudaMemcpyAsync\[1\].*700.*cleanup synchronization.*701", + ): + cuda.copy_d2h( + [PointerCopy(dst=100, src=1000, nbytes=16), PointerCopy(dst=200, src=2000, nbytes=32)], + stream=77, + ) + + assert api.sync_calls == [77] + + +def test_empty_copy_batch_does_not_call_cuda() -> None: + api = FakeCudaBinding() + cuda = CudaTransferRuntime(api=api) + + cuda.copy_d2h([], stream=77) + cuda.copy_h2d(iter(()), stream=77) + + assert api.copy_calls == [] + assert api.sync_calls == [] diff --git a/tests/test_examples_smoke.py b/tests/test_examples_smoke.py index 47cbb32..a0bdb02 100644 --- a/tests/test_examples_smoke.py +++ b/tests/test_examples_smoke.py @@ -19,3 +19,21 @@ def test_approximate_quickstart_lifecycle(): def test_exact_prefix_quickstart_trie(): _run("examples/exact_prefix_reuse/quickstart_trie.py") + + +def test_fluxon_example_uses_sync_plan_store() -> None: + source = (ROOT / "examples/exact_prefix_reuse/e2e_telefuser_lingbot.py").read_text() + fluxon_block = source.split('if args.store == "fluxon"', 1)[1].split( + 'if args.store == "localdisk"', 1 + )[0] + + assert "PlanTensorTierStore(FluxonKVStore" in fluxon_block + assert "async_put=True" not in fluxon_block + + +def test_plan_tier_store_has_no_cpu_staging_or_legacy_tensor_api() -> None: + source = (ROOT / "cacheseek/stores/tier.py").read_text() + plan_block = source.split("class PlanTensorTierStore", 1)[1] + + for forbidden in ("torch.frombuffer", ":spec", "put_tensor", "get_tensor", "queue.Queue"): + assert forbidden not in plan_block diff --git a/tests/test_fluxon_cuda_integration.py b/tests/test_fluxon_cuda_integration.py new file mode 100644 index 0000000..850ad99 --- /dev/null +++ b/tests/test_fluxon_cuda_integration.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Opt-in integration coverage for direct CUDA and Fluxon Plan transfers.""" + +from __future__ import annotations + +import ctypes +import os +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +import torch + +from cacheseek.reuse.exact_prefix.telefuser_lingbot import _RingKVWindow +from cacheseek.stores import BlobHandle, PlanTensorTierStore, Tier +from cacheseek.stores.cuda_transfer import CudaTransferRuntime, PointerCopy +from cacheseek.stores.fluxon import FluxonKVStore + + +class CountingCudaApi: + def __init__(self, api) -> None: + self._api = api + self.sync_calls = 0 + + def __getattr__(self, name): + return getattr(self._api, name) + + def cudaStreamSynchronize(self, stream): + self.sync_calls += 1 + return self._api.cudaStreamSynchronize(stream) + + +@pytest.mark.cuda +def test_registered_host_cuda_round_trip() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + import cuda.bindings.runtime as cuda_api + + source = torch.arange(257, device="cuda", dtype=torch.bfloat16) + target = torch.empty_like(source) + nbytes = source.numel() * source.element_size() + host_buffer = ctypes.create_string_buffer(nbytes) + host_ptr = ctypes.addressof(host_buffer) + generation = os.getpid() + key = (host_ptr, nbytes, generation) + counting_api = CountingCudaApi(cuda_api) + cuda = CudaTransferRuntime(api=counting_api) + stream = cuda.current_stream(source.device) + + try: + cuda.register_fluxon_segments( + [ + { + "write_ptr": host_ptr, + "read_ptr": host_ptr, + "len": nbytes, + "generation": generation, + } + ] + ) + cuda.copy_d2h( + [PointerCopy(dst=host_ptr, src=source.data_ptr(), nbytes=nbytes)], + stream=stream, + ) + assert counting_api.sync_calls == 1 + cuda.copy_h2d( + [PointerCopy(dst=target.data_ptr(), src=host_ptr, nbytes=nbytes)], + stream=stream, + ) + assert counting_api.sync_calls == 2 + assert torch.equal(source, target) + finally: + result = cuda_api.cudaHostUnregister(host_ptr) + CudaTransferRuntime._registered.discard(key) + assert CudaTransferRuntime._code(result) == int(cuda_api.cudaError_t.cudaSuccess) + + +@pytest.mark.cuda +@pytest.mark.fluxon +def test_real_fluxon_plan_gpu_round_trip() -> None: + config = os.environ.get("CACHESEEK_FLUXON_CONFIG") + if not config: + pytest.skip("CACHESEEK_FLUXON_CONFIG is not set") + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + raw = FluxonKVStore(config_path=config) + store = PlanTensorTierStore(raw) + prefix = f"cacheseek-cuda-it-{os.getpid()}-{uuid4().hex}" + kv_locator = prefix + ":kv" + latent_locator = prefix + ":lat" + source_k = torch.arange(24, device="cuda", dtype=torch.bfloat16).reshape(1, 3, 2, 4) + source_v = (source_k + 100).contiguous() + source_latent = torch.arange(48, device="cuda", dtype=torch.bfloat16).reshape( + 1, 2, 3, 2, 4 + ) + ready: list[bool] = [] + try: + store.put_skeleton(latent_locator, source_latent) + store.put_async( + kv_locator, + [(source_k, source_v)], + tier=Tier.FLUXON_DRAM, + on_ready=lambda: ready.append(True), + ) + assert ready == [True] + + restored_latent = torch.empty_like(source_latent) + assert store.materialize_skeletons([(latent_locator, restored_latent)]) + assert torch.equal(restored_latent, source_latent) + + runtime = SimpleNamespace( + frame_tokens=1, + chunk_size=3, + kv_sink_size=0, + self_kv_cache=[ + { + "k": torch.zeros_like(source_k), + "v": torch.zeros_like(source_v), + "global_end_index": 0, + "local_end_index": 0, + } + ], + ) + node = SimpleNamespace( + depth=0, + blob=BlobHandle( + tier=Tier.FLUXON_DRAM, + locator=kv_locator, + nbytes=source_k.nbytes + source_v.nbytes, + n_layers=1, + ready=True, + ), + ) + assert store.materialize_path([node], _RingKVWindow(runtime), depth=0) + assert torch.equal(runtime.self_kv_cache[0]["k"], source_k) + assert torch.equal(runtime.self_kv_cache[0]["v"], source_v) + assert runtime.self_kv_cache[0]["local_end_index"] == 3 + assert runtime.self_kv_cache[0]["global_end_index"] == 3 + finally: + for key_to_remove in ( + latent_locator, + kv_locator + ":L0:k", + kv_locator + ":L0:v", + ): + raw.remove(key_to_remove) diff --git a/tests/test_fluxon_import_identity.py b/tests/test_fluxon_import_identity.py new file mode 100644 index 0000000..8dc3424 --- /dev/null +++ b/tests/test_fluxon_import_identity.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Fluxon package resolution and loaded-extension identity tests.""" + +from __future__ import annotations + +import hashlib +import importlib.machinery +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +from loguru import logger + +import cacheseek.stores.fluxon as fluxon_module +from cacheseek.stores.fluxon import FluxonKVStore + + +@contextmanager +def _captured_messages() -> Iterator[list[str]]: + messages: list[str] = [] + sink_id = logger.add(lambda message: messages.append(message.record["message"])) + try: + yield messages + finally: + logger.remove(sink_id) + + +def _clear_loaded_fluxon_pyo3(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in list(sys.modules): + if module_name == "fluxon_pyo3" or module_name.startswith("fluxon_pyo3."): + monkeypatch.delitem(sys.modules, module_name, raising=False) + + +def _native_module(module_name: str, path: Path) -> ModuleType: + module = ModuleType(module_name) + module.__file__ = str(path) + return module + + +def test_fluxon_import_uses_normal_resolver_without_mutating_sys_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + package_init = tmp_path / "selected-fluxon" / "fluxon_py" / "__init__.py" + package_init.parent.mkdir(parents=True) + package_init.write_text('__version__ = "0.2.1"\n', encoding="utf-8") + api = ModuleType("fluxon_py") + api.__file__ = str(package_init) + api.__version__ = "0.2.1" # type: ignore[attr-defined] + import_calls: list[str] = [] + + def fake_import_module(module_name: str) -> ModuleType: + import_calls.append(module_name) + return api + + monkeypatch.setattr(fluxon_module.importlib, "import_module", fake_import_module) + original_sys_path = list(sys.path) + + with _captured_messages() as messages: + imported = object.__new__(FluxonKVStore)._import_fluxon_layer() + + assert imported is api + assert import_calls == ["fluxon_py"] + assert sys.path == original_sys_path + assert any( + "Fluxon Python package identity" in message + and "module=fluxon_py" in message + and "version=0.2.1" in message + and f"path={package_init.resolve()}" in message + for message in messages + ) + + +def test_loaded_native_extensions_log_resolved_paths_and_sha256( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_loaded_fluxon_pyo3(monkeypatch) + suffix = importlib.machinery.EXTENSION_SUFFIXES[0] + identities = [ + ("fluxon_pyo3", tmp_path / f"direct{suffix}", b"direct extension"), + ( + "fluxon_pyo3.fluxon_pyo3", + tmp_path / f"nested{suffix}", + b"nested extension", + ), + ] + for module_name, extension, payload in identities: + extension.write_bytes(payload) + monkeypatch.setitem( + sys.modules, + module_name, + _native_module(module_name, extension), + ) + fluxon_module._file_sha256.cache_clear() + + with _captured_messages() as messages: + FluxonKVStore._log_fluxon_native_identity() + + for module_name, extension, payload in identities: + expected_sha256 = hashlib.sha256(payload).hexdigest() + assert any( + "Fluxon native extension identity" in message + and f"module={module_name}" in message + and f"path={extension.resolve()}" in message + and f"sha256={expected_sha256}" in message + for message in messages + ) + + +def test_identity_hash_failure_does_not_mask_new_store_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _clear_loaded_fluxon_pyo3(monkeypatch) + suffix = importlib.machinery.EXTENSION_SUFFIXES[0] + extension = tmp_path / f"missing_fluxon_pyo3{suffix}" + module_name = "fluxon_pyo3.fluxon_pyo3" + + def fail_new_store(_config: object) -> None: + monkeypatch.setitem( + sys.modules, + module_name, + _native_module(module_name, extension), + ) + raise RuntimeError("backend boom") + + api = SimpleNamespace( + FluxonKvClientConfig=SimpleNamespace(from_file=lambda _path: object()), + new_store=fail_new_store, + ) + monkeypatch.setattr(FluxonKVStore, "_import_fluxon_layer", lambda _self: api) + fluxon_module._file_sha256.cache_clear() + + with _captured_messages() as messages, pytest.raises(RuntimeError, match="backend boom"): + FluxonKVStore(config_path="/tmp/fluxon.yaml") + + assert any( + f"module={module_name}" in message + and f"path={extension.resolve()}" in message + and "sha256=" in message + and "err_type=FileNotFoundError" in message + for message in messages + ) diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py new file mode 100644 index 0000000..c5ea69a --- /dev/null +++ b/tests/test_fluxon_plan_tier_store.py @@ -0,0 +1,669 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Pointer-first PlanTensorTierStore tests with fake Fluxon and CUDA runtimes.""" + +from __future__ import annotations + +import ctypes +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import pytest + +from cacheseek.reuse.exact_prefix import NamespaceForest, WorldKVConfig, WorldKVManager +from cacheseek.stores import BlobHandle, PlanTensorTierStore, Tier +from cacheseek.stores.cuda_transfer import PointerCopy +from cacheseek.stores.fluxon import FluxonKVStore + +MAGIC = 0x4658_504C_414E_5631 + + +@dataclass +class FakeGetResult: + transferable_len: int + + +@dataclass +class FakeGetHandle: + result: FakeGetResult + keys: list[str] + closed: bool = False + + +class FakeFuture: + def __init__(self, *, ok: bool = True, ret_codes: list[int] | None = None) -> None: + self.ok = ok + self.ret_codes = [] if ret_codes is None else ret_codes + + def wait(self): + return self + + def is_ok(self): + return self.ok + + def unwrap(self, msg: str = ""): + return self.ret_codes + + def unwrap_error(self, msg: str = ""): + return RuntimeError("fake commit failed") + + +class FakeFluxonPlanStore: + def __init__(self, values: dict[str, bytearray] | None = None) -> None: + self.values: dict[str, bytearray] = {} if values is None else values + self.plan_blobs: dict[int, object] = {} + self.started_puts: list[list[str]] = [] + self.committed: list[int] = [] + self.aborted: list[int] = [] + self.released: list[int] = [] + self.cancelled = 0 + self.partial_transferable_len: int | None = None + self.fail_commit_at: int | None = None + self.commit_ret_codes: list[int] | None = None + self.events: list[str] = [] + self.bytes_put_calls = 0 + self.bytes_get_calls = 0 + self.tensor_put_calls = 0 + self.tensor_get_calls = 0 + self.fail_decode = False + self.removed: list[str] = [] + self.fail_remove_key: str | None = None + self.segments = [ + { + "write_ptr": 1000, + "read_ptr": 2000, + "len": 4096, + "generation": 7, + "segment_label": "cpu:0", + "node_id": 0, + } + ] + + def _make_plan(self, ptrs: list[int]) -> int: + blob = (ctypes.c_uint64 * (2 + len(ptrs)))() + blob[0] = MAGIC + blob[1] = len(ptrs) + for index, ptr in enumerate(ptrs): + blob[2 + index] = ptr + plan_ptr = ctypes.addressof(blob) + self.plan_blobs[plan_ptr] = blob + return plan_ptr + + def decode_plan_ptr(self, plan_ptr: int, expected_count: int) -> list[int]: + if self.fail_decode: + raise ValueError("bad plan magic") + header = (ctypes.c_uint64 * 2).from_address(plan_ptr) + if int(header[0]) != MAGIC: + raise ValueError("bad plan magic") + if int(header[1]) != expected_count: + raise ValueError("bad plan pointer count") + words = (ctypes.c_uint64 * (2 + expected_count)).from_address(plan_ptr) + return [int(words[2 + index]) for index in range(expected_count)] + + def wait_local_segments_ready(self) -> list[dict[str, Any]]: + self.events.append("segments") + return self.segments + + def local_fast_put_start(self, keys: list[str], value_len: int, opts=None) -> int: + del opts + self.started_puts.append(list(keys)) + self.events.append("start:" + ",".join(keys)) + ptrs = [] + for key in keys: + buffer = bytearray(value_len) + self.values[key] = buffer + ptrs.append(ctypes.addressof((ctypes.c_ubyte * value_len).from_buffer(buffer))) + return self._make_plan(ptrs) + + def local_fast_put_commit(self, plan_ptr: int): + self.committed.append(plan_ptr) + self.events.append(f"commit:{plan_ptr}") + commit_index = len(self.committed) + return FakeFuture( + ok=commit_index != self.fail_commit_at, + ret_codes=self.commit_ret_codes, + ) + + def put_abort(self, plan_ptr: int) -> None: + self.aborted.append(plan_ptr) + self.events.append(f"abort:{plan_ptr}") + + def get_start( + self, + keys: list[str], + *, + prefix_best_effort: bool = True, + atomic_group_lens: list[int] | None = None, + ) -> FakeGetHandle: + del prefix_best_effort, atomic_group_lens + self.events.append("get_start") + if self.partial_transferable_len is not None: + transferable_len = self.partial_transferable_len + else: + transferable_len = 0 + for key in keys: + if key not in self.values: + break + transferable_len += 1 + return FakeGetHandle(FakeGetResult(transferable_len), list(keys)) + + def get_transfer(self, handle: FakeGetHandle) -> int: + handle.closed = True + self.events.append("get_transfer") + keys = handle.keys[: handle.result.transferable_len] + ptrs = [ + ctypes.addressof((ctypes.c_ubyte * len(self.values[key])).from_buffer(self.values[key])) + for key in keys + ] + return self._make_plan(ptrs) + + def cancel_get_transfer(self, handle: FakeGetHandle) -> None: + handle.closed = True + self.cancelled += 1 + self.events.append("cancel") + + def release_views(self, plan_ptr: int) -> None: + self.released.append(plan_ptr) + self.events.append("release") + + def put(self, key: str, value: bytes) -> None: + self.bytes_put_calls += 1 + self.values[key] = bytearray(value) + + def get(self, key: str) -> bytes | None: + self.bytes_get_calls += 1 + value = self.values.get(key) + return None if value is None else bytes(value) + + def put_tensor(self, key: str, tensor: Any) -> None: + del key, tensor + self.tensor_put_calls += 1 + + def get_tensor(self, key: str, **kwargs: Any) -> None: + del key, kwargs + self.tensor_get_calls += 1 + return None + + def remove(self, key: str) -> None: + self.removed.append(key) + self.events.append(f"remove:{key}") + if key == self.fail_remove_key: + raise RuntimeError(f"remove failed: {key}") + self.values.pop(key, None) + + +class FakeCudaTensor: + def __init__( + self, + data: bytes, + *, + shape: tuple[int, ...] | None = None, + dtype: str = "fake.uint8", + device: str = "cuda:0", + is_cuda: bool = True, + contiguous: bool = True, + ) -> None: + self._buffer = ctypes.create_string_buffer(data, len(data)) + self.shape = (len(data),) if shape is None else shape + self.dtype = dtype + self.device = device + self.is_cuda = is_cuda + self._contiguous = contiguous + self._nbytes = len(data) + + def is_contiguous(self) -> bool: + return self._contiguous + + def data_ptr(self) -> int: + return ctypes.addressof(self._buffer) + + def numel(self) -> int: + return self._nbytes + + def element_size(self) -> int: + return 1 + + def bytes(self) -> bytes: + return ctypes.string_at(self.data_ptr(), self._nbytes) + + +class FakeCudaTransfer: + def __init__(self, events: list[str] | None = None) -> None: + self.events = [] if events is None else events + self.registered: list[list[dict[str, Any]]] = [] + self.d2h_batches: list[tuple[list[Any], int]] = [] + self.h2d_batches: list[tuple[list[Any], int]] = [] + self.fail_d2h = False + self.fail_h2d = False + + def current_stream(self, device: Any) -> int: + assert device == "cuda:0" + return 77 + + def register_fluxon_segments(self, segments) -> None: + self.registered.append(list(segments)) + + def copy_d2h(self, copies, *, stream: int) -> None: + batch = list(copies) + self.d2h_batches.append((batch, stream)) + self.events.append("cuda_d2h_batch") + if self.fail_d2h: + raise RuntimeError("copy exploded") + for copy in batch: + ctypes.memmove(copy.dst, copy.src, copy.nbytes) + + def copy_h2d(self, copies, *, stream: int) -> None: + batch = list(copies) + self.h2d_batches.append((batch, stream)) + self.events.append("cuda_h2d_batch") + if self.fail_h2d: + raise RuntimeError("h2d exploded") + for copy in batch: + ctypes.memmove(copy.dst, copy.src, copy.nbytes) + + +class FakePointerWindow: + def __init__(self, events: list[str]) -> None: + self.events = events + self.resume_depth: int | None = None + self.targets = ctypes.create_string_buffer(128) + self.seen_blobs: list[tuple[int, list[tuple[int, tuple[int, int]]], int]] = [] + + @property + def cuda_device(self) -> str: + return "cuda:0" + + def chunk_value_nbytes(self, layer: int) -> int: + assert layer == 0 + return 4 + + def build_seed_copies(self, layer, blobs, depth): + self.seen_blobs.append((layer, list(blobs), depth)) + copies = [] + for index, (_chunk_depth, (source_k, source_v)) in enumerate(blobs): + copies.extend( + [ + PointerCopy( + dst=ctypes.addressof(self.targets) + index * 8, + src=source_k, + nbytes=4, + ), + PointerCopy( + dst=ctypes.addressof(self.targets) + index * 8 + 4, + src=source_v, + nbytes=4, + ), + ] + ) + return copies + + def set_resume_depth(self, depth: int) -> None: + self.resume_depth = depth + self.events.append("resume") + + +class FakeRawFluxonStore: + def __init__(self) -> None: + self.segments = [ + { + "write_ptr": 1000, + "read_ptr": 2000, + "len": 4096, + "generation": 7, + "segment_label": "cpu:0", + "node_id": 0, + } + ] + + def wait_local_segments_ready(self): + return self.segments + + def put(self, key, value): + raise NotImplementedError + + def get(self, key): + raise NotImplementedError + + def remove(self, key): + raise NotImplementedError + + +def _store(fake: FakeFluxonPlanStore | None = None) -> tuple[FakeFluxonPlanStore, FakeCudaTransfer, Any]: + backend = FakeFluxonPlanStore() if fake is None else fake + cuda = FakeCudaTransfer(backend.events) + return backend, cuda, PlanTensorTierStore(backend, cuda_runtime=cuda) + + +def test_fluxon_wrapper_exposes_segment_mapping_and_not_tensor_api() -> None: + raw = FakeRawFluxonStore() + wrapped = FluxonKVStore(store=raw) + + assert wrapped.wait_local_segments_ready() == raw.segments + assert not hasattr(wrapped, "put_tensor") + assert not hasattr(wrapped, "get_tensor") + + +def test_fluxon_wrapper_keeps_bytes_api() -> None: + wrapped = FluxonKVStore(store=FakeRawFluxonStore()) + + for name in ("put", "get", "remove", "list_keys"): + assert callable(getattr(wrapped, name)) + + +def test_kv_put_uses_one_d2h_batch_and_commits_after_copy() -> None: + fake, cuda, store = _store() + ready: list[bool] = [] + payload = [ + (FakeCudaTensor(b"kkkk"), FakeCudaTensor(b"vvvv")), + (FakeCudaTensor(b"KKKK"), FakeCudaTensor(b"VVVV")), + ] + + store.put_async( + "chunk", + payload, + tier=Tier.FLUXON_DRAM, + on_ready=lambda: (fake.events.append("ready"), ready.append(True)), + ) + + assert len(cuda.d2h_batches) == 1 + copies, stream = cuda.d2h_batches[0] + assert stream == 77 + assert [copy.nbytes for copy in copies] == [4, 4, 4, 4] + assert fake.started_puts == [["chunk:L0:k", "chunk:L0:v", "chunk:L1:k", "chunk:L1:v"]] + assert bytes(fake.values["chunk:L0:k"]) == b"kkkk" + assert bytes(fake.values["chunk:L1:v"]) == b"VVVV" + assert fake.events.index("cuda_d2h_batch") < fake.events.index(f"commit:{fake.committed[0]}") + assert fake.events[-1] == "ready" + assert ready == [True] + + +def test_skeleton_put_uses_plan_d2h_not_put_or_put_tensor() -> None: + fake, cuda, store = _store() + + store.put_skeleton("latent", FakeCudaTensor(b"latent-bytes")) + + assert bytes(fake.values["latent"]) == b"latent-bytes" + assert len(cuda.d2h_batches) == 1 + assert fake.bytes_put_calls == 0 + assert fake.tensor_put_calls == 0 + + +def test_free_removes_only_layer_kv_keys_and_is_idempotent() -> None: + fake, _cuda, store = _store() + fake.values.update( + { + "chunk:L0:k": bytearray(b"kkkk"), + "chunk:L0:v": bytearray(b"vvvv"), + "chunk:L1:k": bytearray(b"KKKK"), + "chunk:L1:v": bytearray(b"VVVV"), + "latent": bytearray(b"keep"), + } + ) + handle = BlobHandle(Tier.FLUXON_DRAM, "chunk", 16, 2, ready=True) + + store.free(handle) + store.free(handle) + + expected = ["chunk:L0:k", "chunk:L0:v", "chunk:L1:k", "chunk:L1:v"] + assert fake.removed == expected * 2 + assert fake.values == {"latent": bytearray(b"keep")} + + +def test_manager_keeps_blob_when_plan_free_fails() -> None: + fake, _cuda, store = _store() + fake.fail_remove_key = "chunk:L0:v" + forest = NamespaceForest() + namespace = forest.get_or_create_namespace(b"r" * 32, b"c" * 32) + node = forest.commit(namespace, namespace.root, "action", b"n" * 32, 0) + handle = BlobHandle(Tier.FLUXON_DRAM, "chunk", 8, 1, ready=True) + node.blob = handle + manager = WorldKVManager( + forest, + store, + WorldKVConfig(window_chunks=1, sink_chunks=0, break_even_k=1), + ) + + with pytest.raises(RuntimeError, match="remove failed"): + manager.evict_blob(node) + + assert node.blob is handle + + +def test_copy_failure_aborts_every_uncommitted_plan() -> None: + fake, cuda, store = _store() + cuda.fail_d2h = True + ready: list[bool] = [] + payload = [(FakeCudaTensor(b"kkkk"), FakeCudaTensor(b"vvvvvv"))] + + with pytest.raises(RuntimeError, match="copy exploded"): + store.put_async( + "chunk", payload, tier=Tier.FLUXON_DRAM, on_ready=lambda: ready.append(True) + ) + + assert len(fake.started_puts) == 2 + assert len(fake.aborted) == 2 + assert fake.committed == [] + assert ready == [] + + +def test_commit_failure_does_not_publish_ready_blob() -> None: + fake = FakeFluxonPlanStore() + fake.fail_commit_at = 1 + fake, _cuda, store = _store(fake) + handle = BlobHandle(Tier.FLUXON_DRAM, "chunk", 8, 1, ready=False) + + with pytest.raises(RuntimeError, match="local_fast_put_commit failed"): + store.put_async( + handle.locator, + [(FakeCudaTensor(b"kkkk"), FakeCudaTensor(b"vvvv"))], + tier=handle.tier, + on_ready=lambda: setattr(handle, "ready", True), + ) + + assert not handle.ready + + +@pytest.mark.parametrize( + ("tensor", "message"), + [ + (FakeCudaTensor(b"data", is_cuda=False), "CUDA"), + (FakeCudaTensor(b"data", contiguous=False), "contiguous"), + ], +) +def test_non_cuda_or_noncontiguous_tensor_is_rejected(tensor, message: str) -> None: + fake, _cuda, store = _store() + + with pytest.raises(ValueError, match=message): + store.put_async("chunk", [(tensor, FakeCudaTensor(b"vvvv"))], tier=Tier.FLUXON_DRAM) + + assert fake.started_puts == [] + + +def test_mixed_nbytes_starts_multiple_plans_but_syncs_one_batch() -> None: + fake, cuda, store = _store() + + store.put_async( + "chunk", + [(FakeCudaTensor(b"kkkk"), FakeCudaTensor(b"vvvvvv"))], + tier=Tier.FLUXON_DRAM, + on_ready=lambda: fake.events.append("ready"), + ) + + assert fake.started_puts == [["chunk:L0:k"], ["chunk:L0:v"]] + assert len(cuda.d2h_batches) == 1 + assert [copy.nbytes for copy in cuda.d2h_batches[0][0]] == [4, 6] + assert [event.split(":", 1)[0] for event in fake.events if event != "segments"] == [ + "start", + "start", + "cuda_d2h_batch", + "commit", + "commit", + "ready", + ] + + +def test_fluxon_segments_are_registered_once() -> None: + _fake, cuda, store = _store() + + store.put_skeleton("latent-a", FakeCudaTensor(b"aaaa")) + store.put_skeleton("latent-b", FakeCudaTensor(b"bbbb")) + + assert cuda.registered == [[ + { + "write_ptr": 1000, + "read_ptr": 2000, + "len": 4096, + "generation": 7, + "segment_label": "cpu:0", + "node_id": 0, + } + ]] + + +def _path(fake: FakeFluxonPlanStore, count: int = 2): + nodes = [] + for depth in range(count): + locator = f"chunk-{depth}" + fake.values[f"{locator}:L0:k"] = bytearray(f"k{depth}!!".encode()) + fake.values[f"{locator}:L0:v"] = bytearray(f"v{depth}!!".encode()) + nodes.append( + SimpleNamespace( + depth=depth, + blob=BlobHandle( + tier=Tier.FLUXON_DRAM, + locator=locator, + nbytes=8, + n_layers=1, + ready=True, + ), + ) + ) + return nodes + + +def test_kv_partial_prefix_cancels_without_h2d_or_index_publish() -> None: + fake, cuda, store = _store() + path = _path(fake) + fake.partial_transferable_len = 2 + window = FakePointerWindow(fake.events) + + assert not store.materialize_path(path, window, depth=1) + + assert fake.cancelled == 1 + assert fake.released == [] + assert cuda.h2d_batches == [] + assert window.resume_depth is None + + +def test_kv_h2d_failure_releases_views_without_index_publish() -> None: + fake, cuda, store = _store() + path = _path(fake) + window = FakePointerWindow(fake.events) + cuda.fail_h2d = True + + with pytest.raises(RuntimeError, match="h2d exploded"): + store.materialize_path(path, window, depth=1) + + assert len(fake.released) == 1 + assert fake.cancelled == 0 + assert window.resume_depth is None + + +def test_kv_full_hit_syncs_then_publishes_indices_then_releases() -> None: + fake, cuda, store = _store() + path = _path(fake) + window = FakePointerWindow(fake.events) + + assert store.materialize_path(path, window, depth=1) + + assert len(cuda.h2d_batches) == 1 + assert window.resume_depth == 1 + assert bytes(window.targets[:16]) == b"k0!!v0!!k1!!v1!!" + assert fake.events[-3:] == ["cuda_h2d_batch", "resume", "release"] + + +def test_kv_plan_decode_failure_releases_views() -> None: + fake, _cuda, store = _store() + path = _path(fake) + window = FakePointerWindow(fake.events) + fake.fail_decode = True + + with pytest.raises(ValueError, match="bad plan magic"): + store.materialize_path(path, window, depth=1) + + assert len(fake.released) == 1 + assert window.resume_depth is None + + +def test_kv_layout_mismatch_is_rejected_before_get_start() -> None: + fake, _cuda, store = _store() + path = _path(fake, count=1) + path[0].blob.nbytes = 7 + window = FakePointerWindow(fake.events) + + with pytest.raises(ValueError, match="layout mismatch"): + store.materialize_path(path, window, depth=0) + + assert "get_start" not in fake.events + + +def test_latents_restore_in_one_h2d_batch() -> None: + fake, cuda, store = _store() + fake.values.update( + { + "latent-0": bytearray(b"aaaa"), + "latent-1": bytearray(b"bbbb"), + "latent-2": bytearray(b"cccc"), + } + ) + targets = [FakeCudaTensor(b"\0" * 4) for _ in range(3)] + + assert store.materialize_skeletons( + [(f"latent-{index}", target) for index, target in enumerate(targets)] + ) + + assert len(cuda.h2d_batches) == 1 + assert [target.bytes() for target in targets] == [b"aaaa", b"bbbb", b"cccc"] + assert len(fake.released) == 1 + + +def test_latent_partial_miss_cancels_without_h2d() -> None: + fake, cuda, store = _store() + fake.values["latent-0"] = bytearray(b"aaaa") + targets = [FakeCudaTensor(b"\0" * 4), FakeCudaTensor(b"\0" * 4)] + + assert not store.materialize_skeletons( + [("latent-0", targets[0]), ("latent-1", targets[1])] + ) + + assert fake.cancelled == 1 + assert fake.released == [] + assert cuda.h2d_batches == [] + + +def test_manager_dispatches_plan_materialize_capability() -> None: + class CapabilityStore: + def __init__(self) -> None: + self.calls = [] + + def materialize_path(self, path, window, *, depth: int) -> bool: + self.calls.append((list(path), window, depth)) + return True + + forest = NamespaceForest() + namespace = forest.get_or_create_namespace(b"r" * 32, b"c" * 32) + node = forest.commit(namespace, namespace.root, "action", b"n" * 32, 0) + node.blob = BlobHandle(Tier.FLUXON_DRAM, "chunk", 8, 1, ready=True) + store = CapabilityStore() + manager = WorldKVManager( + forest, + store, + WorldKVConfig(window_chunks=1, sink_chunks=0, break_even_k=1), + ) + window = object() + + assert manager.materialize(node, window) + + assert store.calls == [([node], window, 0)] + assert node.ref_count == 0 diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index bf5310d..284cdd5 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -14,23 +14,28 @@ """ from __future__ import annotations -import sys -from pathlib import Path +import ctypes from types import SimpleNamespace +import pytest import torch -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from cacheseek.reuse.exact_prefix import ( # noqa: E402 +from cacheseek.reuse.exact_prefix import ( InMemoryTierStore, NamespaceForest, WorldKVManager, + build_action_chain, ) -from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( # noqa: E402 +from cacheseek.reuse.exact_prefix.keys import canonical_json_bytes +from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( LingBotWorldKVBinding, + _RingKVWindow, + chunk_action_keys, make_rolling_config, + session_root_hash, ) +from cacheseek.reuse.exact_prefix.trie import Skeleton +from cacheseek.stores import BlobHandle, Tier # Tiny geometry: chunk=3 frames, frame_tokens=2; window L=7 frames including # sink S=3 frames => rolling starts at the 3rd chunk. @@ -67,8 +72,8 @@ def make_runtime(seed: int, actions: list[int]) -> SimpleNamespace: control_chunks=[torch.full((1, 4), float(a)) for a in actions], self_kv_cache=[ { - "k": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM)), - "v": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM)), + "k": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM), dtype=torch.bfloat16), + "v": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM), dtype=torch.bfloat16), "global_end_index": 0, "local_end_index": 0, } @@ -163,6 +168,241 @@ def make_stack(): return forest, WorldKVManager(forest, InMemoryTierStore(), cfg) +class CudaLikeTensor: + """Expose a CPU tensor pointer through the CUDA tensor protocol used by the store.""" + + def __init__(self, tensor: torch.Tensor) -> None: + self.tensor = tensor + self.is_cuda = True + self.device = "cuda:0" + + @property + def shape(self): + return self.tensor.shape + + @property + def dtype(self): + return self.tensor.dtype + + def is_contiguous(self) -> bool: + return self.tensor.is_contiguous() + + def element_size(self) -> int: + return self.tensor.element_size() + + def data_ptr(self) -> int: + return self.tensor.data_ptr() + + +class RecordingPlanTierStore: + def __init__(self) -> None: + self.events: list[str] = [] + self.latent_result = True + self.kv_result = True + self.raise_latent = False + self.targets: list[tuple[str, torch.Tensor]] = [] + self.saved_latent: torch.Tensor | None = None + self.saved_payload = None + + def materialize_skeletons(self, targets) -> bool: + self.events.append("materialize_skeletons") + self.targets = list(targets) + if self.raise_latent: + raise RuntimeError("latent copy exploded") + if not self.latent_result: + return False + for depth, (_locator, target) in enumerate(self.targets): + target.fill_(depth + 1) + return True + + def materialize_path(self, path, window, *, depth: int) -> bool: + del path + self.events.append("materialize_path") + if self.kv_result: + window.set_resume_depth(depth) + return self.kv_result + + def put_skeleton(self, locator: str, latent: torch.Tensor) -> None: + del locator + self.events.append("put_skeleton") + self.saved_latent = latent + + def put_async(self, locator, payload, *, tier, on_ready=None) -> None: + del locator, tier + self.events.append("put_async") + self.saved_payload = list(payload) + if on_ready is not None: + on_ready() + + +class DeviceTrackingTensor: + """CPU-backed tensor double that records its logical device transfers.""" + + def __init__( + self, + tensor: torch.Tensor, + *, + device: str, + events: list[tuple[str, str]], + ) -> None: + self.tensor = tensor + self.device = device + self.events = events + + @property + def shape(self): + return self.tensor.shape + + @property + def dtype(self): + return self.tensor.dtype + + @property + def nbytes(self) -> int: + return self.tensor.nbytes + + def __getitem__(self, item): + return DeviceTrackingTensor(self.tensor[item], device=self.device, events=self.events) + + def detach(self): + self.events.append(("detach", self.device)) + return self + + def to(self, *, device: str, copy: bool = False): + self.events.append(("to", device)) + tensor = self.tensor.clone() if copy else self.tensor + return DeviceTrackingTensor(tensor, device=device, events=self.events) + + def data_ptr(self) -> int: + return self.tensor.data_ptr() + + +class RecordingGenericTierStore: + def __init__(self) -> None: + self.saved_latent: DeviceTrackingTensor | None = None + self.saved_payload: list[tuple[DeviceTrackingTensor, DeviceTrackingTensor]] | None = None + + def put_skeleton(self, locator: str, latent: DeviceTrackingTensor) -> None: + del locator + self.saved_latent = latent + + def put_async(self, locator, payload, *, tier, on_ready=None) -> None: + del locator, tier + self.saved_payload = list(payload) + if on_ready is not None: + on_ready() + + +def _make_plan_hit( + runtime: SimpleNamespace, + session: SimpleNamespace, + store: RecordingPlanTierStore, + *, + prefix_len: int = 2, +) -> LingBotWorldKVBinding: + forest = NamespaceForest() + cfg = make_rolling_config( + local_attn_size=LOCAL_ATTN, + sink_size=SINK, + chunk_size=CHUNK_FRAMES, + ) + manager = WorldKVManager(forest, store, cfg) + root = session_root_hash(session, model_fingerprint=b"lingbot-world-fast", runtime=runtime) + namespace = forest.get_or_create_namespace(root, root) + actions = chunk_action_keys(runtime) + chain = build_action_chain(root, [canonical_json_bytes(action) for action in actions]) + parent = namespace.root + chunk_nbytes = sum( + 2 * tensor["k"][:, :CT].numel() * tensor["k"].element_size() + for tensor in runtime.self_kv_cache + ) + for depth in range(prefix_len): + node = forest.commit(namespace, parent, actions[depth], chain[depth], depth) + node.skeleton = Skeleton(latent_locator=f"latent-{depth}") + node.blob = BlobHandle( + tier=Tier.FLUXON_DRAM, + locator=f"chunk-{depth}", + nbytes=chunk_nbytes, + n_layers=len(runtime.self_kv_cache), + ready=True, + ) + parent = node + return LingBotWorldKVBinding(manager, forest, ingest_enabled=False) + + +def _pointer_runtime() -> tuple[SimpleNamespace, list[tuple[torch.Tensor, torch.Tensor]]]: + batch = 2 + chunks = [] + for chunk in range(N_CHUNKS): + values = torch.arange(batch * CT * HEADS * HEAD_DIM, dtype=torch.int32).reshape( + batch, CT, HEADS, HEAD_DIM + ) + k = (values + chunk * 10_000).contiguous() + chunks.append((k, (k + 1000).contiguous())) + target_k = torch.zeros((batch, KV_TOKENS, HEADS, HEAD_DIM), dtype=torch.int32) + target_v = torch.zeros_like(target_k) + runtime = SimpleNamespace( + frame_tokens=FRAME_TOKENS, + chunk_size=CHUNK_FRAMES, + kv_sink_size=SINK, + self_kv_cache=[ + { + "k": CudaLikeTensor(target_k), + "v": CudaLikeTensor(target_v), + "global_end_index": 0, + "local_end_index": 0, + } + ], + ) + return runtime, chunks + + +@pytest.mark.parametrize("chunks_to_restore", [1, 2, 3, 4]) +def test_ring_pointer_copies_match_tensor_reference(chunks_to_restore: int) -> None: + depth = chunks_to_restore - 1 + pointer_runtime, chunks = _pointer_runtime() + pointer_window = _RingKVWindow(pointer_runtime) + + reference_k = torch.zeros((2, KV_TOKENS, HEADS, HEAD_DIM), dtype=torch.int32) + reference_v = torch.zeros_like(reference_k) + reference_runtime = SimpleNamespace( + frame_tokens=FRAME_TOKENS, + chunk_size=CHUNK_FRAMES, + kv_sink_size=SINK, + self_kv_cache=[ + { + "k": reference_k, + "v": reference_v, + "global_end_index": 0, + "local_end_index": 0, + } + ], + ) + reference_window = _RingKVWindow(reference_runtime) + tensor_blobs = [(chunk, chunks[chunk]) for chunk in range(chunks_to_restore)] + reference_window.seed_layer(0, tensor_blobs, depth) + reference_window.set_resume_depth(depth) + + pointer_blobs = [ + (chunk, (chunks[chunk][0].data_ptr(), chunks[chunk][1].data_ptr())) + for chunk in range(chunks_to_restore) + ] + copies = pointer_window.build_seed_copies(0, pointer_blobs, depth) + for copy in copies: + ctypes.memmove(copy.dst, copy.src, copy.nbytes) + pointer_window.set_resume_depth(depth) + + pointer_k = pointer_runtime.self_kv_cache[0]["k"].tensor + pointer_v = pointer_runtime.self_kv_cache[0]["v"].tensor + local_end = reference_runtime.self_kv_cache[0]["local_end_index"] + assert pointer_runtime.self_kv_cache[0]["local_end_index"] == local_end + assert pointer_runtime.self_kv_cache[0]["global_end_index"] == ( + reference_runtime.self_kv_cache[0]["global_end_index"] + ) + assert torch.equal(pointer_k[:, :local_end], reference_k[:, :local_end]) + assert torch.equal(pointer_v[:, :local_end], reference_v[:, :local_end]) + + # ---------------------------------------------------------------------- tests def test_warm_ring_equals_cold_ring_all_resume_points(): """K=1..4: the warm-restored ring is bit-equal to a cold run's ring just @@ -224,6 +464,151 @@ def test_namespace_isolation_by_seed(): assert b.last_fast_forward == 0 +def test_session_root_hash_isolates_runtime_layout() -> None: + session = make_session(42) + runtime_a = make_runtime(42, [1, 2, 3, 4]) + runtime_b = make_runtime(42, [1, 2, 3, 4]) + runtime_b.self_kv_cache.append( + { + "k": torch.zeros_like(runtime_b.self_kv_cache[0]["k"]), + "v": torch.zeros_like(runtime_b.self_kv_cache[0]["v"]), + "global_end_index": 0, + "local_end_index": 0, + } + ) + + root_a = session_root_hash(session, model_fingerprint=b"weights", runtime=runtime_a) + root_b = session_root_hash(session, model_fingerprint=b"weights", runtime=runtime_b) + + assert root_a != root_b + + +def test_latent_partial_miss_cancels_before_kv_materialize() -> None: + runtime = make_runtime(42, [1, 2, 9, 9]) + store = RecordingPlanTierStore() + store.latent_result = False + binding = _make_plan_hit(runtime, make_session(42), store) + + binding.on_runtime_created(runtime, make_session(42)) + + assert store.events == ["materialize_skeletons"] + assert binding.last_fast_forward == 0 + assert runtime.world_kv_cached_latents == {} + assert all(kv["local_end_index"] == kv["global_end_index"] == 0 for kv in runtime.self_kv_cache) + + +def test_latent_copy_failure_leaves_runtime_indices_zero() -> None: + runtime = make_runtime(42, [1, 2, 9, 9]) + store = RecordingPlanTierStore() + store.raise_latent = True + binding = _make_plan_hit(runtime, make_session(42), store) + + with pytest.raises(RuntimeError, match="latent copy exploded"): + binding.on_runtime_created(runtime, make_session(42)) + + assert store.events == ["materialize_skeletons"] + assert binding.last_fast_forward == 0 + assert runtime.world_kv_cached_latents == {} + assert all(kv["local_end_index"] == kv["global_end_index"] == 0 for kv in runtime.self_kv_cache) + + +def test_cached_latents_are_private_targets_in_depth_order() -> None: + runtime = make_runtime(42, [1, 2, 9, 9]) + store = RecordingPlanTierStore() + binding = _make_plan_hit(runtime, make_session(42), store) + + binding.on_runtime_created(runtime, make_session(42)) + + assert store.events == ["materialize_skeletons", "materialize_path"] + assert binding.last_fast_forward == 2 + assert list(runtime.world_kv_cached_latents) == [0, 1] + assert [locator for locator, _target in store.targets] == ["latent-0", "latent-1"] + assert list(runtime.world_kv_cached_latents.values()) == [ + target for _locator, target in store.targets + ] + assert all(target.dtype == runtime.self_kv_cache[0]["k"].dtype for _, target in store.targets) + + +@pytest.mark.parametrize( + "denoised", + [ + torch.empty((1,), dtype=torch.bfloat16), + torch.empty((1, 2, CHUNK_FRAMES, 2, 2), dtype=torch.float32), + ], +) +def test_latent_layout_mismatch_is_rejected_before_plan_start(denoised: torch.Tensor) -> None: + runtime = make_runtime(42, [1, 2, 3, 4]) + session = make_session(42) + store = RecordingPlanTierStore() + binding = _make_plan_hit(runtime, session, store, prefix_len=0) + binding.ingest_enabled = True + binding.on_runtime_created(runtime, session) + + with pytest.raises(ValueError, match="latent.*(shape|dtype) mismatch"): + binding.on_chunk_finalized(runtime, 0, denoised) + + assert "put_skeleton" not in store.events + assert "put_async" not in store.events + + +def test_plan_write_uses_device_views_without_cloning() -> None: + runtime = make_runtime(42, [1, 2, 3, 4]) + session = make_session(42) + store = RecordingPlanTierStore() + binding = _make_plan_hit(runtime, session, store, prefix_len=0) + binding.ingest_enabled = True + binding.on_runtime_created(runtime, session) + denoised = torch.empty_like(runtime.noise_chunks[0], dtype=runtime.self_kv_cache[0]["k"].dtype) + for kv in runtime.self_kv_cache: + kv["local_end_index"] = CT + + binding.on_chunk_finalized(runtime, 0, denoised) + + assert store.saved_latent is not None + assert store.saved_latent.data_ptr() == denoised.data_ptr() + assert store.saved_payload is not None + for kv, (saved_k, saved_v) in zip(runtime.self_kv_cache, store.saved_payload, strict=True): + assert saved_k.data_ptr() == kv["k"].data_ptr() + assert saved_v.data_ptr() == kv["v"].data_ptr() + + +def test_generic_write_stages_owned_cpu_tensors() -> None: + runtime = make_runtime(42, [1, 2, 3, 4]) + session = make_session(42) + store = RecordingGenericTierStore() + binding = _make_plan_hit(runtime, session, store, prefix_len=0) + binding.ingest_enabled = True + binding.on_runtime_created(runtime, session) + + events: list[tuple[str, str]] = [] + source_pairs = [] + for kv in runtime.self_kv_cache: + source_k = DeviceTrackingTensor(kv["k"], device="cuda:0", events=events) + source_v = DeviceTrackingTensor(kv["v"], device="cuda:0", events=events) + kv["k"], kv["v"] = source_k, source_v + kv["local_end_index"] = CT + source_pairs.append((source_k, source_v)) + denoised = DeviceTrackingTensor( + torch.empty_like(runtime.noise_chunks[0], dtype=source_pairs[0][0].dtype), + device="cuda:0", + events=events, + ) + + binding.on_chunk_finalized(runtime, 0, denoised) + + assert store.saved_latent is not None + assert store.saved_latent.device == "cpu" + assert store.saved_latent.data_ptr() != denoised.data_ptr() + assert store.saved_payload is not None + for (source_k, source_v), (saved_k, saved_v) in zip( + source_pairs, store.saved_payload, strict=True + ): + assert saved_k.device == saved_v.device == "cpu" + assert saved_k.data_ptr() != source_k.data_ptr() + assert saved_v.data_ptr() != source_v.data_ptr() + assert events.count(("to", "cpu")) == 2 * N_LAYERS + 1 + + if __name__ == "__main__": for fn in [v for k, v in sorted(globals().items()) if k.startswith("test_")]: fn()