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

[yJader]

Summary

This PR replaces the legacy per-tensor Fluxon path used by exact-prefix reuse with a synchronous Plan-pointer data path that transfers LingBot KV and skeleton latents directly between CUDA tensors and Fluxon registered-host segments.

The new path:

• writes GPU KV/latent bytes through local_fast_put_start() + batched CUDA D2H copies + local_fast_put_commit();
• restores cached latents and the active KV window through get_start() / get_transfer() + batched CUDA H2D copies;
• keeps Fluxon Plan handles and pointers inside explicit commit/abort/cancel/release lifetimes;
• publishes cached runtime state only after the required transfers complete;
• removes the legacy Fluxon put_tensor/get_tensor path, CPU tensor staging, and :spec sidecars from exact-prefix traffic;
• resolves fluxon_py only through the configured Python environment and logs the loaded package/native-extension identity;
• preserves the generic bytes API for approximate reuse and the generic tensor adapter's owned CPU snapshots for LocalDisk;
• reclaims a Plan blob's Fluxon K/V values when the existing explicit WorldKVManager.evict_blob() entry point is called.

Related Fluxon API work: Tele-AI/Fluxon#42.

Motivation

The previous Fluxon exact-prefix path operated one tensor at a time and materialized CPU-owned tensors around the DLPack interface:

GPU KV
-> CPU tensor staging
-> Fluxon put_tensor

Fluxon get_tensor
-> holder/DLPack access
-> CPU clone/materialization
-> runtime GPU KV

That model had three problems for LingBot prefix reuse:

1. it inserted CPU tensor staging into both the write and read hot paths;
2. it issued storage operations per tensor instead of restoring the complete active window as one ordered request;
3. it did not expose the pointer lifetime needed to copy directly into the TeleFuser physical KV ring.

Fluxon's Plan API exposes writable or readable value pointers with explicit lifecycle methods. This PR uses those pointers as the boundary between CacheSeek's runtime layout logic and Fluxon's storage ownership.

Architecture

LingBotWorldKVBinding
  - runtime hooks, latent-first restore, CUDA KV views
              |
              v
WorldKVManager
  - trie/window selection and BlobHandle.ready
              |
              v
PlanTensorTierStore
  - key ordering, size grouping, Plan lifecycle, copy batches
       |                         |
       v                         v
FluxonKVStore             CudaTransferRuntime
  - Plan API wrapper        - host registration
  - plan_ptr decode         - H2D/D2H + stream sync

The main implementation is split across:

• cacheseek/stores/cuda_transfer.py: injectable CUDA runtime facade and PointerCopy descriptors;
• cacheseek/stores/fluxon.py: Fluxon bytes adapter plus thin wrappers for the Plan API;
• cacheseek/stores/tier.py: synchronous PlanTensorTierStore;
• cacheseek/reuse/exact_prefix/manager.py: materialize_path() capability dispatch;
• cacheseek/reuse/exact_prefix/telefuser_lingbot.py: direct ring-copy descriptors, private latent targets, and layout isolation;
• examples/exact_prefix_reuse/e2e_telefuser_lingbot.py: Fluxon example selection now uses PlanTensorTierStore.

Write path

LingBotWorldKVBinding.on_chunk_finalized() now passes detached CUDA views for the finalized chunk instead of cloning them to CPU.

The binding keeps the backend split explicit: the Fluxon Plan path receives those CUDA views, while LocalDisk and other generic tensor stores receive detached, caller-owned CPU clones. This preserves the generic adapter's copy-before-overwrite contract and keeps CUDA tensors away from LocalDiskTensorStore.put_tensor().numpy().

For each skeleton latent or KV chunk, PlanTensorTierStore:

1. validates that all inputs are contiguous CUDA tensors on one device;
2. flattens KV payloads to {locator}:L{layer}:k/v keys;
3. groups entries by nbytes, because one local_fast_put_start(keys, value_len) applies a single value size to every key in that Plan;
4. starts all required put Plans and decodes their destination pointers;
5. submits one D2H descriptor batch on the current PyTorch CUDA stream;
6. synchronizes the stream once;
7. commits each Plan and checks both the future Result and per-key return codes;
8. calls on_ready() only after every commit succeeds.

BlobHandle.ready therefore remains false until all KV values are committed. A failure before a Plan is committed calls put_abort() for every live uncommitted Plan. If a CUDA batch was only partially submitted, its stream is synchronized before those aborts begin.

Although the manager method is still named put_async(), the Fluxon Plan implementation is intentionally synchronous. Its flush() method is a no-op; LocalDisk retains the existing asynchronous generic tensor adapter.

Read path

The hit path restores private skeleton latents before mutating the KV ring:

private latent targets
-> latent Plan GET + H2D
-> KV-window Plan GET + H2D
-> publish KV indices
-> publish cached latents / trie cursor / fast-forward depth

For skeletons, the binding allocates private GPU targets using the current runtime shape, dtype, and device. materialize_skeletons() fetches all latent keys in one Plan request and copies them in one H2D batch. A missing latent cancels the handle and leaves runtime-visible restore state untouched.

For KV, WorldKVManager.materialize() detects the store's materialize_path() capability and passes the complete sink + recent-window path to PlanTensorTierStore.

KV keys are ordered chunk-major:

chunk0 L0:k, chunk0 L0:v, ..., chunk0 LN:v,
chunk1 L0:k, chunk1 L0:v, ..., chunk1 LN:v,
...

Each chunk is one atomic group:

atomic_group_lens = [n_layers * 2] * len(path)

The current implementation consumes only a complete requested window. A partial transferable prefix is cancelled and treated as a cache miss because it is not necessarily equivalent to the sink + recent-window path of a shallower trie node.

_RingKVWindow.build_seed_copies() translates each Fluxon chunk pointer into contiguous copy runs in the physical TeleFuser ring, including rolled layouts with sink frames. KV indices are updated only after the H2D batch synchronizes successfully. Read views are then released in finally.

Fluxon Plan and CUDA lifetimes

Both get_transfer() and local_fast_put_start() return a pointer to a small u64 table:

word[0] = 0x4658_504c_414e_5631
word[1] = value pointer count
word[2..] = value pointers in key order

FluxonKVStore.decode_plan_ptr() validates the address, magic, and pointer count before exposing any value pointer.

Lifecycle rules enforced by the adapter:

• a get handle is consumed by either get_transfer() or cancel_get_transfer();
• every copy descriptor is validated before the first cudaMemcpyAsync() submission;
• if a later cudaMemcpyAsync() fails after earlier copies were submitted, the stream is synchronized before any Plan is aborted or its views are released;
• a read Plan remains alive until the H2D stream is synchronized, then release_views() runs;
• a put Plan remains alive until D2H completes, then it is committed or aborted;
• a submitted put Plan is not aborted again, even if its future fails;
• BlobHandle.ready and runtime resume indices are never published before their corresponding transfer/commit succeeds.

Before the first transfer, wait_local_segments_ready() returns Fluxon owner mappings. The CUDA facade registers both write and read mappings, handles read-only fallback and already-registered results, and deduplicates registrations by (ptr, len, generation).

Layout isolation

Plan values contain raw tensor bytes and no shape/dtype metadata. The read path derives byte lengths and offsets from the active runtime and validates them against BlobHandle.nbytes/n_layers.

The exact-prefix namespace fingerprint now includes:

storage_format = "fluxon-plan-python-cuda-v1"
n_layers
kv_shape / kv_dtype
latent_shape / latent_dtype
chunk_size / frame_tokens

Changing the storage format or runtime layout therefore creates a different namespace. Legacy DLPack/CPU-Plan values and incompatible layouts become cold misses instead of being reinterpreted.

Compatibility and scope

• Exact-prefix Fluxon traffic uses PlanTensorTierStore.
• Approximate reuse continues to use KVStore.put/get(bytes).
• FluxonKVStore.put/get/remove/list_keys remain available.
• TensorKVStore, TensorStoreTierStore, and LocalDiskTensorStore remain available for generic tensor backends.
• InMemory and LocalDisk continue through the existing per-layer get_layer()/seed_layer() fallback; generic writeback explicitly stages owned CPU tensors.
• WorldKVManager.evict_blob(node) now removes the node's {locator}:L{layer}:k/v Plan values while retaining its lightweight skeleton.
• FluxonKVStore uses normal Python package resolution and never scans parent directories or modifies sys.path.
• Production deployments should pin a Fluxon wheel/package; source development must use an editable install or explicit PYTHONPATH.
• Startup logs the resolved fluxon_py version/path and each loaded fluxon_pyo3 extension path/SHA256.
• The new cuda-bindings>=12.6 dependency is isolated in the fluxon optional extra.
• The Fluxon/CUDA integration tests are opt-in through the new cuda and fluxon pytest markers.

Correctness coverage

The unit and fake Fluxon/CUDA tests cover:

• Plan magic/count validation;
• write/read segment registration and registration deduplication;
• read-only registration fallback;
• D2H/H2D direction, descriptor order, and one stream sync per batch;
• whole-batch descriptor validation before submission and cleanup synchronization after a partially submitted batch fails;
• same-size and mixed-size put groups;
• sync-before-commit and ready-after-commit ordering;
• copy, decode, commit, and nonzero return-code failures;
• put abort, get cancel, and view release paths;
• full KV-window hits and partial-prefix cancellation;
• latent-first restore and no runtime index publication on latent failure;
• full-length, rolled, and sink + recent physical-ring reconstruction;
• runtime-layout namespace isolation;
• no CPU staging or legacy tensor API use in the Plan tier, plus owned CPU snapshots for generic tensor stores;
• explicit Plan K/V reclamation, skeleton retention, idempotent removal, and blob retention after a removal failure;
• standard Fluxon package resolution without sys.path mutation;
• package version/path and direct or nested native-extension path/SHA256 logging;
• backend-init/hash failure diagnostics;
• Fluxon example selection of the synchronous Plan store.

Tests

Run on the reviewed CacheSeek worktree:

export CUDA_MPS_PIPE_DIRECTORY=/tmp/cuda-mps-disabled-$USER
export PYTHONPATH=.
python -m pytest -q \
  tests/test_cuda_transfer.py \
  tests/test_fluxon_plan_tier_store.py \
  tests/test_world_kv_telefuser_binding.py \
  tests/test_world_kv_flow.py \
  tests/test_fluxon_import_identity.py \
  tests/test_examples_smoke.py
python -m ruff check .

Results:

• pytest: 56 passed, 1 environment warning (asyncio_mode was unknown because the active pytest environment did not load pytest-asyncio);
• Ruff: all checks passed.

Real-device validation from the recorded benchmark stack:

Test	Result
Registered-host CUDA round trip	1 passed / 2.776 s
Real Fluxon Plan KV + latent GPU round trip	1 passed / 101.880 s
Plan A/B/C/D isolated-service samples	3/3 passed all correctness checks

The A/B/C/D correctness gate verifies:

• A and D are cold references;
• B hits all three chunks and every output frame hash matches A;
• C hits the first two chunks and every output frame hash matches D.

Performance evidence

The recorded benchmark used an NVIDIA H100 80 GB, LingBot-World-Fast, 37 output frames, three chunks, a two-chunk fork prefix, fixed seed, and an independent fresh 64 GiB Fluxon service for each sample.

Median of three samples:

Case	Legacy tensor API	Plan API	Legacy / Plan	Wall reduction
A: cold + first write	18.806 s	70.025 s	0.27x	-272.4%
B: full hit	9.116 s	1.881 s	4.85x	79.4%
C: 2/3 prefix hit	12.521 s	3.820 s	3.28x	69.5%
D: cold fork reference	18.639 s	7.229 s	2.58x	61.2%

Instrumented transfer results:

Path	Payload	Time	Throughput	Scope
Plan H2D	17.855 GiB	0.348 s	51.28 GiB/s	CUDA pointer copies
Plan D2H	24.998 GiB	0.542 s	46.14 GiB/s	CUDA pointer copies
Legacy get	17.855 GiB	6.604 s	2.70 GiB/s	complete get_tensor path
Legacy put	24.998 GiB	29.241 s	0.85 GiB/s	complete put_tensor path

The full-hit and prefix-hit results show the read/materialization benefit. The cold first-write result is a known regression: a fresh service currently pays roughly 50–60 seconds of one-time registered-host segment setup. D2H copy itself accounted for only about 0.23 seconds of the first 10.71 GiB write.

These figures are indicative rather than release claims: the tested Fluxon revision, upstream PR head, TeleFuser revision, allocator settings, and closed-SDK binary were not all identical. See the validation notes below.

Known limitations and follow-ups

1. First-use registration. The fresh-service first write pays a 50–60 second one-time segment setup cost. This should move to an explicit service warm-up/initialization phase and gain native stage timers.
2. Capacity/reclaim policy. PlanTensorTierStore.free() now lets an explicit WorldKVManager.evict_blob(node) call reclaim that node's Fluxon K/V values without deleting its skeleton. CacheSeek still has no automatic capacity accounting or LRU victim selection, so a workload that continuously writes without explicitly evicting eligible nodes can still exhaust a 64 GiB owner and make local_fast_put_start() time out.
3. Synchronous request path. The implementation synchronizes each latent/KV copy batch and waits for commit futures. It does not yet overlap copy, compute, or remote placement.
4. Trie/Fluxon state consistency. Fluxon may evict committed values under capacity pressure while CacheSeek retains the corresponding ready blob handles in its trie. A stale trie match can therefore produce an incomplete atomic-group-aligned transferable prefix. The read safely falls back to cold computation, but CacheSeek does not currently invalidate the stale trie state or retry the longest shorter materializable prefix.
5. Orphan values after partial commit. If one size-group Plan commits and a later Plan fails, the trie node remains unpublished but the earlier values may remain in Fluxon until backend cleanup.
6. TeleFuser compatibility scope. This PR only integrates CacheSeek with the new Fluxon Plan API. Adapting the integration and E2E example to the latest TeleFuser runtime is out of scope and remains follow-up work.