From ea52bf0a460ab1bcb84c1913406a9ca8cf83ec49 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:02:00 +0000 Subject: [PATCH 01/11] feat(stores): add Python CUDA transfer facade --- cacheseek/stores/cuda_transfer.py | 111 ++++++++++++++++++++++++++ pyproject.toml | 1 + tests/test_cuda_transfer.py | 128 ++++++++++++++++++++++++++++++ 3 files changed, 240 insertions(+) create mode 100644 cacheseek/stores/cuda_transfer.py create mode 100644 tests/test_cuda_transfer.py diff --git a/cacheseek/stores/cuda_transfer.py b/cacheseek/stores/cuda_transfer.py new file mode 100644 index 0000000..ed13e4b --- /dev/null +++ b/cacheseek/stores/cuda_transfer.py @@ -0,0 +1,111 @@ +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}") + self._check( + self._api.cudaMemcpyAsync(copy.dst, copy.src, copy.nbytes, kind, stream), + f"cudaMemcpyAsync[{index}]", + ) + 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/pyproject.toml b/pyproject.toml index 50e22e7..319bc6b 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]"] diff --git a/tests/test_cuda_transfer.py b/tests/test_cuda_transfer.py new file mode 100644 index 0000000..0765f7c --- /dev/null +++ b/tests/test_cuda_transfer.py @@ -0,0 +1,128 @@ +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]] = [] + + 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: + self.sync_calls.append(int(stream)) + 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_does_not_sync() -> 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 == [] + + +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 == [] From 3a2a841711bac84e3ce7937d79ce00c3b02e2991 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:04:21 +0000 Subject: [PATCH 02/11] refactor(fluxon): remove legacy tensor API --- cacheseek/stores/__init__.py | 8 +- cacheseek/stores/base.py | 12 +- cacheseek/stores/fluxon.py | 233 +++++++++------- tests/test_fluxon_plan_tier_store.py | 398 +++++++++++++++++++++++++++ 4 files changed, 536 insertions(+), 115 deletions(-) create mode 100644 tests/test_fluxon_plan_tier_store.py 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/fluxon.py b/cacheseek/stores/fluxon.py index 05ad4e1..4720885 100644 --- a/cacheseek/stores/fluxon.py +++ b/cacheseek/stores/fluxon.py @@ -15,9 +15,11 @@ ## 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. Prefer local source under ``fluxon_new`` or ``Fluxon`` at the TeleFuser + workspace root. +2. Fall back to the historical ``fluxon/pylib_src`` layout, then an installed + ``fluxon_py`` package. +3. If all fail, raise ``ImportError`` at the call site. The module itself always imports as long as ``__init__`` is not actually invoked (see ``__init__.py``). @@ -51,6 +53,7 @@ from __future__ import annotations +import ctypes from typing import Any from loguru import logger @@ -60,6 +63,7 @@ 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, @@ -186,114 +190,124 @@ def list_keys(self) -> list[str]: 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 +333,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 "" @@ -352,7 +366,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,9 +406,9 @@ 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. + Prefer local source under ``fluxon_new`` or ``Fluxon`` at the TeleFuser + workspace root, then the historical ``fluxon/pylib_src`` path, falling + back to an installed ``fluxon_py`` package. Raise ImportError if all fail. """ load_errors: list[str] = [] @@ -403,12 +417,19 @@ def _import_fluxon_layer(self): 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(): + here = _Path(__file__).resolve() + candidates = [] + for parent in here.parents: + candidates.extend( + [ + parent / "fluxon_new", + parent / "Fluxon", + parent / "fluxon" / "pylib_src", + ] + ) + for fluxon_root in candidates: + if not fluxon_root.is_dir(): + continue fluxon_path = str(fluxon_root) if fluxon_path not in sys.path: sys.path.insert(0, fluxon_path) @@ -435,5 +456,5 @@ def _import_fluxon_layer(self): ) raise ImportError( f"Fluxon Python API not available (fluxon_py): {detail}. " - "Install fluxon_py or place the source at /fluxon/pylib_src/." + "Install fluxon_py or place the source at /fluxon_new." ) diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py new file mode 100644 index 0000000..c5b0abf --- /dev/null +++ b/tests/test_fluxon_plan_tier_store.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""PlanTensorTierStore tests using an in-process fake Fluxon plan backend.""" + +from __future__ import annotations + +import ctypes +from dataclasses import dataclass + +import pytest +import torch + +from cacheseek.reuse.exact_prefix import ( + NamespaceForest, + PlanTensorTierStore, + WorldKVConfig, + WorldKVManager, + build_action_chain, + root_hash, +) +from cacheseek.reuse.exact_prefix.keys import canonical_json_bytes, sha256 +from cacheseek.stores.fluxon import FluxonKVStore + +MAGIC = 0x4658_504C_414E_5631 +N_LAYERS = 2 + + +@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=""): + return self.ret_codes + + def unwrap_error(self, msg=""): + 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 = False + self.commit_ret_codes: list[int] | None = None + + def _make_plan(self, ptrs: list[int]) -> int: + blob = (ctypes.c_uint64 * (2 + len(ptrs)))() + blob[0] = MAGIC + blob[1] = len(ptrs) + for i, ptr in enumerate(ptrs): + blob[2 + i] = 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]: + header = (ctypes.c_uint64 * 2).from_address(plan_ptr) + assert int(header[0]) == MAGIC + assert int(header[1]) == expected_count + words = (ctypes.c_uint64 * (2 + expected_count)).from_address(plan_ptr) + return [int(words[2 + i]) for i in range(expected_count)] + + def local_fast_put_start(self, keys: list[str], value_len: int, opts=None) -> int: + self.started_puts.append(list(keys)) + ptrs = [] + for key in keys: + buf = bytearray(value_len) + self.values[key] = buf + ptrs.append(ctypes.addressof((ctypes.c_ubyte * value_len).from_buffer(buf))) + return self._make_plan(ptrs) + + def local_fast_put_commit(self, plan_ptr: int): + self.committed.append(plan_ptr) + return FakeFuture(ok=not self.fail_commit, ret_codes=self.commit_ret_codes) + + def put_abort(self, plan_ptr: int) -> None: + self.aborted.append(plan_ptr) + + def get_start( + self, + keys: list[str], + *, + prefix_best_effort: bool = True, + atomic_group_lens: list[int] | None = None, + ) -> FakeGetHandle: + 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=transferable_len), list(keys)) + + def get_transfer(self, handle: FakeGetHandle) -> int: + handle.closed = True + 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 + + def release_views(self, plan_ptr: int) -> None: + self.released.append(plan_ptr) + + def put(self, key: str, value: bytes) -> None: + self.values[key] = bytearray(value) + + def get(self, key: str) -> bytes | None: + value = self.values.get(key) + return None if value is None else bytes(value) + + +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 + + +class FakeWindow: + def __init__(self) -> None: + self.seeded: list[tuple[int, list, int]] = [] + self.resume_depth: int | None = None + + def seed_layer(self, layer, blobs, depth): + self.seeded.append((layer, list(blobs), depth)) + + def set_resume_depth(self, depth): + self.resume_depth = depth + + +class RaisingWindow(FakeWindow): + def seed_layer(self, layer, blobs, depth): + raise RuntimeError("seed exploded") + + +def _root() -> bytes: + return root_hash( + image_fp=sha256(b"img", b"A"), + prompt_fp=sha256(b"p", b"A"), + config_blob_hash=sha256(b"cfg", b"{}"), + ) + + +def _payload(nk: bytes) -> list[tuple[torch.Tensor, torch.Tensor]]: + base = int.from_bytes(nk[:2], "little") % 1000 + out = [] + for layer in range(N_LAYERS): + k = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) + base + layer * 100 + out.append((k.contiguous(), (k + 1000).contiguous())) + return out + + +def _drive(mgr, forest, root, actions): + ns = forest.get_or_create_namespace(root, b"cfgblob") + chain = build_action_chain(root, [canonical_json_bytes(a) for a in actions]) + node = ns.root + for i, node_key in enumerate(chain): + node = mgr.ingest(ns, node, actions[i], node_key, i, _payload(node_key), latent=f"x0@{i}") + return chain + + +def _make_stack(fake: FakeFluxonPlanStore): + forest = NamespaceForest() + store = PlanTensorTierStore(fake) + cfg = WorldKVConfig(window_chunks=2, sink_chunks=1, break_even_k=1) + return forest, store, WorldKVManager(forest, store, cfg) + + +def test_fluxon_wrapper_exposes_segment_mapping_and_not_tensor_api(): + 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(): + wrapped = FluxonKVStore(store=FakeRawFluxonStore()) + + for name in ("put", "get", "remove", "list_keys"): + assert callable(getattr(wrapped, name)) + + +def test_plan_tier_writes_with_local_fast_put_and_materializes_with_get_transfer(): + fake = FakeFluxonPlanStore() + forest, store, mgr = _make_stack(fake) + root = _root() + actions = ["a", "b", "c"] + chain = _drive(mgr, forest, root, actions) + + assert len(fake.started_puts) == len(actions) + assert len(fake.committed) == len(actions) + assert not fake.aborted + + win = FakeWindow() + res = mgr.try_fast_forward(root, actions, win) + assert res.start_chunk == 3 + assert win.resume_depth == 2 + assert len(fake.released) == 1 + assert fake.cancelled == 0 + + # window = sink d0 + recent d1,d2, seeded per layer. + assert len(win.seeded) == N_LAYERS + layer0, blobs0, depth0 = win.seeded[0] + assert layer0 == 0 and depth0 == 2 + assert [d for d, _ in blobs0] == [0, 1, 2] + expected = [_payload(chain[i])[0] for i in (0, 1, 2)] + for (_, got), exp in zip(blobs0, expected, strict=True): + assert torch.equal(got[0], exp[0]) + assert torch.equal(got[1], exp[1]) + + +def test_partial_transferable_prefix_cancels_and_falls_back(): + fake = FakeFluxonPlanStore() + forest, store, mgr = _make_stack(fake) + root = _root() + actions = ["a", "b", "c"] + _drive(mgr, forest, root, actions) + fake.partial_transferable_len = 2 + + win = FakeWindow() + res = mgr.try_fast_forward(root, actions, win) + assert res.start_chunk == 0 + assert fake.cancelled == 1 + assert not fake.released + assert win.seeded == [] + + +def test_read_exception_releases_transferred_views(): + fake = FakeFluxonPlanStore() + forest, store, mgr = _make_stack(fake) + root = _root() + actions = ["a", "b", "c"] + _drive(mgr, forest, root, actions) + + with pytest.raises(RuntimeError, match="seed exploded"): + mgr.try_fast_forward(root, actions, RaisingWindow()) + + assert len(fake.released) == 1 + assert fake.cancelled == 0 + + +def test_copy_failure_aborts_uncommitted_plan(monkeypatch): + import cacheseek.stores.tier as tier_mod + + fake = FakeFluxonPlanStore() + forest, store, mgr = _make_stack(fake) + root = _root() + ns = forest.get_or_create_namespace(root, b"cfgblob") + chain = build_action_chain(root, [canonical_json_bytes("a")]) + + def fail_memmove(*args): + raise RuntimeError("copy exploded") + + monkeypatch.setattr(tier_mod.ctypes, "memmove", fail_memmove) + + with pytest.raises(RuntimeError, match="copy exploded"): + mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") + + node = ns.root.children["a"] + assert node.blob is not None + assert not node.blob.ready + assert len(fake.started_puts) == 1 + assert len(fake.aborted) == 1 + assert not fake.committed + + +def test_mixed_tensor_sizes_create_multiple_local_fast_put_plans(): + fake = FakeFluxonPlanStore() + forest, store, mgr = _make_stack(fake) + root = _root() + ns = forest.get_or_create_namespace(root, b"cfgblob") + chain = build_action_chain(root, [canonical_json_bytes("a")]) + payload = [ + ( + torch.arange(12, dtype=torch.float32).reshape(1, 3, 4), + torch.arange(6, dtype=torch.float32).reshape(1, 3, 2), + ) + ] + + mgr.ingest(ns, ns.root, "a", chain[0], 0, payload, latent="x0") + + assert len(fake.started_puts) == 2 + assert sorted(len(group) for group in fake.started_puts) == [1, 1] + assert len(fake.committed) == 2 + assert not fake.aborted + + +def test_commit_failure_does_not_publish_ready_blob(): + fake = FakeFluxonPlanStore() + fake.fail_commit = True + forest, store, mgr = _make_stack(fake) + root = _root() + ns = forest.get_or_create_namespace(root, b"cfgblob") + chain = build_action_chain(root, [canonical_json_bytes("a")]) + + with pytest.raises(RuntimeError, match="local_fast_put_commit failed"): + mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") + + node = ns.root.children["a"] + assert node is not None + assert node.blob is not None + assert not node.blob.ready + + +def test_nonzero_commit_ret_code_does_not_publish_ready_blob(): + fake = FakeFluxonPlanStore() + fake.commit_ret_codes = [0, 7] + forest, store, mgr = _make_stack(fake) + root = _root() + ns = forest.get_or_create_namespace(root, b"cfgblob") + chain = build_action_chain(root, [canonical_json_bytes("a")]) + + with pytest.raises(RuntimeError, match="nonzero ret codes"): + mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") + + node = ns.root.children["a"] + assert node.blob is not None + assert not node.blob.ready + + +def test_plan_tier_loads_specs_from_sidecars_in_new_store_instance(): + backing: dict[str, bytearray] = {} + fake_a = FakeFluxonPlanStore(backing) + forest_a, store_a, mgr_a = _make_stack(fake_a) + root = _root() + actions = ["a", "b", "c"] + _drive(mgr_a, forest_a, root, actions) + snap = forest_a.snapshot() + + fake_b = FakeFluxonPlanStore(backing) + forest_b = NamespaceForest() + forest_b.load_snapshot(snap) + store_b = PlanTensorTierStore(fake_b) + mgr_b = WorldKVManager( + forest_b, + store_b, + WorldKVConfig(window_chunks=2, sink_chunks=1, break_even_k=1), + ) + + win = FakeWindow() + res = mgr_b.try_fast_forward(root, actions, win) + assert res.start_chunk == 3 + assert win.resume_depth == 2 + assert len(fake_b.released) == 1 From f8b5554c1395aa02fc9cb573296d58e6afef3edc Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:09:46 +0000 Subject: [PATCH 03/11] feat(exact-prefix): write Fluxon plans directly from GPU --- cacheseek/stores/tier.py | 172 ++++++++++- tests/test_fluxon_plan_tier_store.py | 444 ++++++++++++++------------- 2 files changed, 401 insertions(+), 215 deletions(-) diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index 2d85263..f7af019 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -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: @@ -172,7 +181,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 +249,164 @@ 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: + """Fluxon plan-API tier store for exact-prefix KV blobs. + + Unlike ``TensorStoreTierStore``, this adapter does not route through + per-key ``put_tensor/get_tensor``. It writes raw tensor bytes into + ``local_fast_put_start`` value pointers and restores prefix hits through + ``get_start/get_transfer`` pointer views. + """ + + 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_path(self, path: Sequence[Any], window: Any, *, depth: int) -> bool: + del path, window, depth + return False + + def free(self, handle: BlobHandle) -> None: + del handle diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py index c5b0abf..978d251 100644 --- a/tests/test_fluxon_plan_tier_store.py +++ b/tests/test_fluxon_plan_tier_store.py @@ -1,28 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project -"""PlanTensorTierStore tests using an in-process fake Fluxon plan backend.""" +"""Pointer-first PlanTensorTierStore tests with fake Fluxon and CUDA runtimes.""" from __future__ import annotations import ctypes from dataclasses import dataclass +from typing import Any import pytest -import torch - -from cacheseek.reuse.exact_prefix import ( - NamespaceForest, - PlanTensorTierStore, - WorldKVConfig, - WorldKVManager, - build_action_chain, - root_hash, -) -from cacheseek.reuse.exact_prefix.keys import canonical_json_bytes, sha256 + +from cacheseek.stores import BlobHandle, PlanTensorTierStore, Tier from cacheseek.stores.fluxon import FluxonKVStore MAGIC = 0x4658_504C_414E_5631 -N_LAYERS = 2 @dataclass @@ -38,7 +29,7 @@ class FakeGetHandle: class FakeFuture: - def __init__(self, ok: bool = True, ret_codes: list[int] | None = None) -> None: + 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 @@ -48,10 +39,10 @@ def wait(self): def is_ok(self): return self.ok - def unwrap(self, msg=""): + def unwrap(self, msg: str = ""): return self.ret_codes - def unwrap_error(self, msg=""): + def unwrap_error(self, msg: str = ""): return RuntimeError("fake commit failed") @@ -65,41 +56,70 @@ def __init__(self, values: dict[str, bytearray] | None = None) -> None: self.released: list[int] = [] self.cancelled = 0 self.partial_transferable_len: int | None = None - self.fail_commit = False + 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.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 i, ptr in enumerate(ptrs): - blob[2 + i] = ptr + 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]: header = (ctypes.c_uint64 * 2).from_address(plan_ptr) - assert int(header[0]) == MAGIC - assert int(header[1]) == expected_count + 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 + i]) for i in range(expected_count)] + 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: - buf = bytearray(value_len) - self.values[key] = buf - ptrs.append(ctypes.addressof((ctypes.c_ubyte * value_len).from_buffer(buf))) + 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) - return FakeFuture(ok=not self.fail_commit, ret_codes=self.commit_ret_codes) + 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, @@ -108,6 +128,7 @@ def get_start( prefix_best_effort: bool = True, atomic_group_lens: list[int] | None = None, ) -> FakeGetHandle: + del prefix_best_effort, atomic_group_lens if self.partial_transferable_len is not None: transferable_len = self.partial_transferable_len else: @@ -116,7 +137,7 @@ def get_start( if key not in self.values: break transferable_len += 1 - return FakeGetHandle(FakeGetResult(transferable_len=transferable_len), list(keys)) + return FakeGetHandle(FakeGetResult(transferable_len), list(keys)) def get_transfer(self, handle: FakeGetHandle) -> int: handle.closed = True @@ -135,12 +156,86 @@ def release_views(self, plan_ptr: int) -> None: self.released.append(plan_ptr) 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 + + +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 + + +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 + + 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)) + for copy in batch: + ctypes.memmove(copy.dst, copy.src, copy.nbytes) + class FakeRawFluxonStore: def __init__(self) -> None: @@ -168,57 +263,13 @@ def remove(self, key): raise NotImplementedError -class FakeWindow: - def __init__(self) -> None: - self.seeded: list[tuple[int, list, int]] = [] - self.resume_depth: int | None = None - - def seed_layer(self, layer, blobs, depth): - self.seeded.append((layer, list(blobs), depth)) - - def set_resume_depth(self, depth): - self.resume_depth = depth - - -class RaisingWindow(FakeWindow): - def seed_layer(self, layer, blobs, depth): - raise RuntimeError("seed exploded") - - -def _root() -> bytes: - return root_hash( - image_fp=sha256(b"img", b"A"), - prompt_fp=sha256(b"p", b"A"), - config_blob_hash=sha256(b"cfg", b"{}"), - ) - - -def _payload(nk: bytes) -> list[tuple[torch.Tensor, torch.Tensor]]: - base = int.from_bytes(nk[:2], "little") % 1000 - out = [] - for layer in range(N_LAYERS): - k = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) + base + layer * 100 - out.append((k.contiguous(), (k + 1000).contiguous())) - return out - - -def _drive(mgr, forest, root, actions): - ns = forest.get_or_create_namespace(root, b"cfgblob") - chain = build_action_chain(root, [canonical_json_bytes(a) for a in actions]) - node = ns.root - for i, node_key in enumerate(chain): - node = mgr.ingest(ns, node, actions[i], node_key, i, _payload(node_key), latent=f"x0@{i}") - return chain - +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 _make_stack(fake: FakeFluxonPlanStore): - forest = NamespaceForest() - store = PlanTensorTierStore(fake) - cfg = WorldKVConfig(window_chunks=2, sink_chunks=1, break_even_k=1) - return forest, store, WorldKVManager(forest, store, cfg) - -def test_fluxon_wrapper_exposes_segment_mapping_and_not_tensor_api(): +def test_fluxon_wrapper_exposes_segment_mapping_and_not_tensor_api() -> None: raw = FakeRawFluxonStore() wrapped = FluxonKVStore(store=raw) @@ -227,172 +278,137 @@ def test_fluxon_wrapper_exposes_segment_mapping_and_not_tensor_api(): assert not hasattr(wrapped, "get_tensor") -def test_fluxon_wrapper_keeps_bytes_api(): +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_plan_tier_writes_with_local_fast_put_and_materializes_with_get_transfer(): - fake = FakeFluxonPlanStore() - forest, store, mgr = _make_stack(fake) - root = _root() - actions = ["a", "b", "c"] - chain = _drive(mgr, forest, root, actions) - - assert len(fake.started_puts) == len(actions) - assert len(fake.committed) == len(actions) - assert not fake.aborted - - win = FakeWindow() - res = mgr.try_fast_forward(root, actions, win) - assert res.start_chunk == 3 - assert win.resume_depth == 2 - assert len(fake.released) == 1 - assert fake.cancelled == 0 - - # window = sink d0 + recent d1,d2, seeded per layer. - assert len(win.seeded) == N_LAYERS - layer0, blobs0, depth0 = win.seeded[0] - assert layer0 == 0 and depth0 == 2 - assert [d for d, _ in blobs0] == [0, 1, 2] - expected = [_payload(chain[i])[0] for i in (0, 1, 2)] - for (_, got), exp in zip(blobs0, expected, strict=True): - assert torch.equal(got[0], exp[0]) - assert torch.equal(got[1], exp[1]) - - -def test_partial_transferable_prefix_cancels_and_falls_back(): - fake = FakeFluxonPlanStore() - forest, store, mgr = _make_stack(fake) - root = _root() - actions = ["a", "b", "c"] - _drive(mgr, forest, root, actions) - fake.partial_transferable_len = 2 - - win = FakeWindow() - res = mgr.try_fast_forward(root, actions, win) - assert res.start_chunk == 0 - assert fake.cancelled == 1 - assert not fake.released - assert win.seeded == [] - +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")), + ] -def test_read_exception_releases_transferred_views(): - fake = FakeFluxonPlanStore() - forest, store, mgr = _make_stack(fake) - root = _root() - actions = ["a", "b", "c"] - _drive(mgr, forest, root, actions) + store.put_async( + "chunk", + payload, + tier=Tier.FLUXON_DRAM, + on_ready=lambda: (fake.events.append("ready"), ready.append(True)), + ) - with pytest.raises(RuntimeError, match="seed exploded"): - mgr.try_fast_forward(root, actions, RaisingWindow()) + 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] - assert len(fake.released) == 1 - assert fake.cancelled == 0 +def test_skeleton_put_uses_plan_d2h_not_put_or_put_tensor() -> None: + fake, cuda, store = _store() -def test_copy_failure_aborts_uncommitted_plan(monkeypatch): - import cacheseek.stores.tier as tier_mod + store.put_skeleton("latent", FakeCudaTensor(b"latent-bytes")) - fake = FakeFluxonPlanStore() - forest, store, mgr = _make_stack(fake) - root = _root() - ns = forest.get_or_create_namespace(root, b"cfgblob") - chain = build_action_chain(root, [canonical_json_bytes("a")]) + 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 fail_memmove(*args): - raise RuntimeError("copy exploded") - monkeypatch.setattr(tier_mod.ctypes, "memmove", fail_memmove) +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"): - mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") + store.put_async( + "chunk", payload, tier=Tier.FLUXON_DRAM, on_ready=lambda: ready.append(True) + ) - node = ns.root.children["a"] - assert node.blob is not None - assert not node.blob.ready - assert len(fake.started_puts) == 1 - assert len(fake.aborted) == 1 - assert not fake.committed + assert len(fake.started_puts) == 2 + assert len(fake.aborted) == 2 + assert fake.committed == [] + assert ready == [] -def test_mixed_tensor_sizes_create_multiple_local_fast_put_plans(): +def test_commit_failure_does_not_publish_ready_blob() -> None: fake = FakeFluxonPlanStore() - forest, store, mgr = _make_stack(fake) - root = _root() - ns = forest.get_or_create_namespace(root, b"cfgblob") - chain = build_action_chain(root, [canonical_json_bytes("a")]) - payload = [ - ( - torch.arange(12, dtype=torch.float32).reshape(1, 3, 4), - torch.arange(6, dtype=torch.float32).reshape(1, 3, 2), + 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), ) - ] - mgr.ingest(ns, ns.root, "a", chain[0], 0, payload, latent="x0") + assert not handle.ready - assert len(fake.started_puts) == 2 - assert sorted(len(group) for group in fake.started_puts) == [1, 1] - assert len(fake.committed) == 2 - assert not fake.aborted +@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() -def test_commit_failure_does_not_publish_ready_blob(): - fake = FakeFluxonPlanStore() - fake.fail_commit = True - forest, store, mgr = _make_stack(fake) - root = _root() - ns = forest.get_or_create_namespace(root, b"cfgblob") - chain = build_action_chain(root, [canonical_json_bytes("a")]) + with pytest.raises(ValueError, match=message): + store.put_async("chunk", [(tensor, FakeCudaTensor(b"vvvv"))], tier=Tier.FLUXON_DRAM) - with pytest.raises(RuntimeError, match="local_fast_put_commit failed"): - mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") + assert fake.started_puts == [] - node = ns.root.children["a"] - assert node is not None - assert node.blob is not None - assert not node.blob.ready +def test_mixed_nbytes_starts_multiple_plans_but_syncs_one_batch() -> None: + fake, cuda, store = _store() -def test_nonzero_commit_ret_code_does_not_publish_ready_blob(): - fake = FakeFluxonPlanStore() - fake.commit_ret_codes = [0, 7] - forest, store, mgr = _make_stack(fake) - root = _root() - ns = forest.get_or_create_namespace(root, b"cfgblob") - chain = build_action_chain(root, [canonical_json_bytes("a")]) - - with pytest.raises(RuntimeError, match="nonzero ret codes"): - mgr.ingest(ns, ns.root, "a", chain[0], 0, _payload(chain[0]), latent="x0") - - node = ns.root.children["a"] - assert node.blob is not None - assert not node.blob.ready - - -def test_plan_tier_loads_specs_from_sidecars_in_new_store_instance(): - backing: dict[str, bytearray] = {} - fake_a = FakeFluxonPlanStore(backing) - forest_a, store_a, mgr_a = _make_stack(fake_a) - root = _root() - actions = ["a", "b", "c"] - _drive(mgr_a, forest_a, root, actions) - snap = forest_a.snapshot() - - fake_b = FakeFluxonPlanStore(backing) - forest_b = NamespaceForest() - forest_b.load_snapshot(snap) - store_b = PlanTensorTierStore(fake_b) - mgr_b = WorldKVManager( - forest_b, - store_b, - WorldKVConfig(window_chunks=2, sink_chunks=1, break_even_k=1), + store.put_async( + "chunk", + [(FakeCudaTensor(b"kkkk"), FakeCudaTensor(b"vvvvvv"))], + tier=Tier.FLUXON_DRAM, + on_ready=lambda: fake.events.append("ready"), ) - win = FakeWindow() - res = mgr_b.try_fast_forward(root, actions, win) - assert res.start_chunk == 3 - assert win.resume_depth == 2 - assert len(fake_b.released) == 1 + 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, + } + ]] From 74861d8451d17a7f75aa05be746164b80ffbdb3d Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:14:23 +0000 Subject: [PATCH 04/11] feat(exact-prefix): restore Fluxon KV directly into GPU ring --- .../reuse/exact_prefix/telefuser_lingbot.py | 94 +++++++++++- cacheseek/stores/tier.py | 76 +++++++++- tests/test_fluxon_plan_tier_store.py | 140 ++++++++++++++++++ tests/test_world_kv_telefuser_binding.py | 110 +++++++++++++- 4 files changed, 411 insertions(+), 9 deletions(-) diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index f66553f..ffd331e 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 @@ -172,6 +173,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 +218,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: diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index f7af019..d894ee7 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -405,8 +405,80 @@ def put_skeleton(self, locator: str, latent: Any) -> None: ) def materialize_path(self, path: Sequence[Any], window: Any, *, depth: int) -> bool: - del path, window, depth - return False + 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: del handle diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py index 978d251..f4e6c83 100644 --- a/tests/test_fluxon_plan_tier_store.py +++ b/tests/test_fluxon_plan_tier_store.py @@ -6,11 +6,13 @@ import ctypes from dataclasses import dataclass +from types import SimpleNamespace from typing import Any import pytest 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 @@ -63,6 +65,7 @@ def __init__(self, values: dict[str, bytearray] | None = None) -> None: self.bytes_get_calls = 0 self.tensor_put_calls = 0 self.tensor_get_calls = 0 + self.fail_decode = False self.segments = [ { "write_ptr": 1000, @@ -85,6 +88,8 @@ def _make_plan(self, ptrs: list[int]) -> int: 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") @@ -129,6 +134,7 @@ def get_start( 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: @@ -141,6 +147,7 @@ def get_start( 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])) @@ -151,9 +158,11 @@ def get_transfer(self, handle: FakeGetHandle) -> int: 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 @@ -213,6 +222,7 @@ def __init__(self, events: list[str] | None = None) -> None: 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" @@ -233,10 +243,53 @@ def copy_d2h(self, copies, *, stream: int) -> None: 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 = [ @@ -412,3 +465,90 @@ def test_fluxon_segments_are_registered_once() -> None: "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 diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index bf5310d..64af034 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -14,21 +14,20 @@ """ 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, ) -from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( # noqa: E402 +from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( LingBotWorldKVBinding, + _RingKVWindow, make_rolling_config, ) @@ -163,6 +162,105 @@ 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() + + +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 From fd1cf1f5ca1c0fd45adb840aa7bb80d6ed828ce9 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:19:22 +0000 Subject: [PATCH 05/11] feat(exact-prefix): restore Fluxon skeletons directly to GPU --- .../reuse/exact_prefix/telefuser_lingbot.py | 108 +++++++--- cacheseek/stores/tier.py | 43 ++++ tests/test_fluxon_plan_tier_store.py | 37 ++++ tests/test_world_kv_telefuser_binding.py | 196 +++++++++++++++++- 4 files changed, 356 insertions(+), 28 deletions(-) diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index ffd331e..92f1697 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -87,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 @@ -114,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) @@ -390,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) @@ -408,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. @@ -462,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: @@ -470,13 +526,13 @@ 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: + k_view = k_view.clone() + v_view = v_view.clone() + payload.append((k_view, v_view)) + latent = denoised.detach() if plan_data_path else denoised.detach().clone() # 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/tier.py b/cacheseek/stores/tier.py index d894ee7..c5b8705 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -404,6 +404,49 @@ def put_skeleton(self, locator: str, latent: Any) -> None: [_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 diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py index f4e6c83..65ace8c 100644 --- a/tests/test_fluxon_plan_tier_store.py +++ b/tests/test_fluxon_plan_tier_store.py @@ -214,6 +214,9 @@ def numel(self) -> int: 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: @@ -552,3 +555,37 @@ def test_kv_layout_mismatch_is_rejected_before_get_start() -> None: 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 == [] diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index 64af034..bcc88c5 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -24,12 +24,18 @@ InMemoryTierStore, NamespaceForest, WorldKVManager, + build_action_chain, ) +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. @@ -66,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, } @@ -188,6 +194,84 @@ 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() + + +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 = [] @@ -322,6 +406,114 @@ 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() + + if __name__ == "__main__": for fn in [v for k, v in sorted(globals().items()) if k.startswith("test_")]: fn() From 742e3868284f28f23334788510629917b0a6a477 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:21:32 +0000 Subject: [PATCH 06/11] docs(exact-prefix): document synchronous Fluxon plan path --- .../reuse/exact_prefix/telefuser_lingbot.py | 4 +-- cacheseek/stores/tier.py | 25 ++++++++----------- .../e2e_telefuser_lingbot.py | 15 +++++------ tests/test_examples_smoke.py | 18 +++++++++++++ 4 files changed, 36 insertions(+), 26 deletions(-) diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index 92f1697..546e09e 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -372,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 diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index c5b8705..395f550 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 @@ -76,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. """ @@ -121,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. @@ -252,12 +249,10 @@ def free(self, handle: BlobHandle) -> None: # noqa: ARG002 -- no-op (reclamati class PlanTensorTierStore: - """Fluxon plan-API tier store for exact-prefix KV blobs. + """Synchronous Fluxon Plan adapter for exact-prefix GPU tensor bytes. - Unlike ``TensorStoreTierStore``, this adapter does not route through - per-key ``put_tensor/get_tensor``. It writes raw tensor bytes into - ``local_fast_put_start`` value pointers and restores prefix hits through - ``get_start/get_transfer`` pointer views. + 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: 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/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 From 48768466af73094123c284c14901b78e4d6dac05 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:23:38 +0000 Subject: [PATCH 07/11] test(fluxon): cover direct CUDA plan round trip --- pyproject.toml | 2 + tests/test_fluxon_cuda_integration.py | 147 ++++++++++++++++++++++++++ 2 files changed, 149 insertions(+) create mode 100644 tests/test_fluxon_cuda_integration.py diff --git a/pyproject.toml b/pyproject.toml index 319bc6b..19f48a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,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_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) From a86269705604c4b58689365a7a3a3c98e526a0db Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 10 Jul 2026 07:25:11 +0000 Subject: [PATCH 08/11] fix(exact-prefix): finalize direct Fluxon CUDA transfer --- cacheseek/reuse/exact_prefix/manager.py | 8 +++++++ tests/test_fluxon_plan_tier_store.py | 28 +++++++++++++++++++++++++ 2 files changed, 36 insertions(+) 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/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py index 65ace8c..c3e0b63 100644 --- a/tests/test_fluxon_plan_tier_store.py +++ b/tests/test_fluxon_plan_tier_store.py @@ -11,6 +11,7 @@ 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 @@ -589,3 +590,30 @@ def test_latent_partial_miss_cancels_without_h2d() -> None: 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 From bc92be0c69258b044643d8191d1eece3eb50e3b9 Mon Sep 17 00:00:00 2001 From: jader Date: Fri, 31 Jul 2026 03:26:47 +0000 Subject: [PATCH 09/11] fix(fluxon): make package import deterministic --- cacheseek/stores/fluxon.py | 168 +++++++++++++++++---------- tests/test_fluxon_import_identity.py | 148 +++++++++++++++++++++++ 2 files changed, 254 insertions(+), 62 deletions(-) create mode 100644 tests/test_fluxon_import_identity.py diff --git a/cacheseek/stores/fluxon.py b/cacheseek/stores/fluxon.py index 4720885..276f302 100644 --- a/cacheseek/stores/fluxon.py +++ b/cacheseek/stores/fluxon.py @@ -7,19 +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_new`` or ``Fluxon`` at the TeleFuser - workspace root. -2. Fall back to the historical ``fluxon/pylib_src`` layout, then an installed - ``fluxon_py`` package. -3. If all 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``). @@ -54,11 +53,37 @@ 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.""" @@ -90,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") @@ -184,8 +213,7 @@ 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 [] @@ -269,7 +297,9 @@ 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.local_fast_put_commit failed plan_ptr={} err={}", plan_ptr, exc) + logger.exception( + "FluxonKV.local_fast_put_commit failed plan_ptr={} err={}", plan_ptr, exc + ) raise RuntimeError( f"FluxonKV.local_fast_put_commit failed plan_ptr={plan_ptr}: {exc}" ) from exc @@ -347,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") @@ -406,55 +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_new`` or ``Fluxon`` at the TeleFuser - workspace root, then the historical ``fluxon/pylib_src`` path, falling - back to an installed ``fluxon_py`` package. Raise ImportError if all 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 - - here = _Path(__file__).resolve() - candidates = [] - for parent in here.parents: - candidates.extend( - [ - parent / "fluxon_new", - parent / "Fluxon", - parent / "fluxon" / "pylib_src", - ] - ) - for fluxon_root in candidates: - if not fluxon_root.is_dir(): - continue - 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}" - ) + 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 - # 2) Fall back to the installed fluxon_py package. - try: - import fluxon_py as api # type: ignore + 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 - return api - except Exception as exc: - load_errors.append( - f"installed fluxon_py import failed type={type(exc).__name__} err={exc}" + @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_new." - ) + 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/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 + ) From 8b0abf16d1325c143d9f53dc86dfaffdbdefc87f Mon Sep 17 00:00:00 2001 From: jader Date: Tue, 4 Aug 2026 06:05:45 +0000 Subject: [PATCH 10/11] fix(exact-prefix): harden Fluxon Plan lifecycle --- .../reuse/exact_prefix/telefuser_lingbot.py | 12 ++- cacheseek/stores/cuda_transfer.py | 26 ++++- cacheseek/stores/tier.py | 9 +- tests/test_cuda_transfer.py | 39 +++++++- tests/test_fluxon_plan_tier_store.py | 50 ++++++++++ tests/test_world_kv_telefuser_binding.py | 95 +++++++++++++++++++ 6 files changed, 221 insertions(+), 10 deletions(-) diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index 546e09e..d852500 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -529,10 +529,16 @@ def on_chunk_finalized( k_view = kv["k"][:, s:e].detach() v_view = kv["v"][:, s:e].detach() if not plan_data_path: - k_view = k_view.clone() - v_view = v_view.clone() + # 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)) - latent = denoised.detach() if plan_data_path else denoised.detach().clone() + 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/cuda_transfer.py b/cacheseek/stores/cuda_transfer.py index ed13e4b..f5488d6 100644 --- a/cacheseek/stores/cuda_transfer.py +++ b/cacheseek/stores/cuda_transfer.py @@ -53,10 +53,28 @@ def _copy_batch(self, copies: Iterable[PointerCopy], *, kind: Any, stream: int) 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}") - self._check( - self._api.cudaMemcpyAsync(copy.dst, copy.src, copy.nbytes, kind, stream), - f"cudaMemcpyAsync[{index}]", - ) + 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: diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index 395f550..e33d3e5 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -519,4 +519,11 @@ def materialize_path(self, path: Sequence[Any], window: Any, *, depth: int) -> b self._ps.release_views(plan_ptr) def free(self, handle: BlobHandle) -> None: - del handle + """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/tests/test_cuda_transfer.py b/tests/test_cuda_transfer.py index 0765f7c..ab143b1 100644 --- a/tests/test_cuda_transfer.py +++ b/tests/test_cuda_transfer.py @@ -20,6 +20,7 @@ def __init__(self) -> None: 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))) @@ -35,8 +36,10 @@ def cudaMemcpyAsync( return self.copy_results.pop(0) return (self.cudaSuccess,) - def cudaStreamSynchronize(self, stream: int) -> int: + 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 @@ -103,7 +106,7 @@ def test_read_only_registration_retries_without_read_only() -> None: ] -def test_memcpy_failure_does_not_sync() -> None: +def test_memcpy_failure_syncs_previously_submitted_work() -> None: api = FakeCudaBinding() api.copy_results = [api.cudaSuccess, (700,)] cuda = CudaTransferRuntime(api=api) @@ -114,9 +117,41 @@ def test_memcpy_failure_does_not_sync() -> None: 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) diff --git a/tests/test_fluxon_plan_tier_store.py b/tests/test_fluxon_plan_tier_store.py index c3e0b63..c5ea69a 100644 --- a/tests/test_fluxon_plan_tier_store.py +++ b/tests/test_fluxon_plan_tier_store.py @@ -67,6 +67,8 @@ def __init__(self, values: dict[str, bytearray] | None = None) -> None: 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, @@ -183,6 +185,13 @@ def get_tensor(self, key: str, **kwargs: Any) -> None: 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__( @@ -380,6 +389,47 @@ def test_skeleton_put_uses_plan_d2h_not_put_or_put_tensor() -> None: 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 diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index bcc88c5..284cdd5 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -235,6 +235,64 @@ def put_async(self, locator, payload, *, tier, on_ready=None) -> 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, @@ -514,6 +572,43 @@ def test_plan_write_uses_device_views_without_cloning() -> None: 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() From b80be68940186a2bd58ec7a706aca7c5b4e54523 Mon Sep 17 00:00:00 2001 From: jader Date: Tue, 4 Aug 2026 06:23:05 +0000 Subject: [PATCH 11/11] fix(exact-prefix): add PlanTensorTierStore import to exact-prefix module --- cacheseek/reuse/exact_prefix/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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", ]