Apple GPU R4: device-resident latent KV cache (cache stops round-tripping) - #38
Conversation
…ping) Final phase of the GPU-resident architecture: c_kv / k_rope live in resident device buffers, so the cache no longer round-trips across decode steps. - runtime.py: DeviceTensor.prefix_view(n_rows) — a non-owning view of the first n_rows rows (offset 0, contiguous), so a growing resident cache exposes its populated window [current_seq, ...] with no copy. - python/tessera/cache/resident_kv.py: ResidentLatentKVCache. Allocates the latent [max_seq, Dl] and rope [max_seq, dr] device buffers ONCE; append writes only the new token in place via the unified-memory .numpy() view (no upload, no realloc); latent_window()/rope_window() return zero-copy prefix-view DeviceTensors that feed device-resident kernels directly. numpy fallback for portability. - cache/__init__.py: export ResidentLatentKVCache. - tests/unit/test_apple_gpu_resident_kv.py: 8 tests — append/window contents, prefill+append, zero-copy aliasing of the resident buffer, resident window feeding a device bmm (R1) with no host copy + numpy cross-check, single allocation across 20 appends, overflow raises, footprint, prefix_view bounds. - docs: plan R4 marked done (contiguous latent cache); device-resident block-paged gather noted as the follow-on. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f0bb26974f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._latent = dt.empty((self.max_seq, self.latent_dim), np.float32) | ||
| self._rope = dt.empty((self.max_seq, self.rope_dim), np.float32) |
There was a problem hiding this comment.
On hosts where the Apple GPU runtime cannot be loaded or built, DeviceTensor.empty() does not reliably return None; it can propagate RuntimeError/CalledProcessError from runtime._load_apple_gpu_runtime() before this constructor reaches the numpy fallback below. That makes ResidentLatentKVCache(...) unusable in exactly the “ABI unavailable” environment the fallback is meant to support, and the new unguarded non-GPU tests hit this path on Linux when the stub build fails. Catch allocation/runtime failures here and treat them the same as None so the portable backing is actually used.
Useful? React with 👍 / 👎.
1 participant
The final phase of the GPU-resident architecture:
c_kv/k_ropelive inresident device buffers, so the KV cache no longer round-trips across decode
steps.
What
DeviceTensor.prefix_view(n_rows)— a non-owning view of the firstn_rowsrows (offset 0, contiguous), so a growing resident cache exposes itspopulated window
[current_seq, …]with no copy.tessera.cache.ResidentLatentKVCache— allocates the latent[max_seq, Dl]+ rope[max_seq, dr]device buffers once.appendwritesonly the new token in place through the unified-memory
.numpy()view — noupload, no reallocation, the GPU sees it immediately.
latent_window()/rope_window()return zero-copy prefix-viewDeviceTensors that feeddevice-resident kernels directly. numpy fallback for portability.
Why it's the last piece
R3 kept weights resident but still re-wrapped the KV window each step. R4
keeps the cache itself resident: across the decode loop, only the new token's
KV is written (a unified-memory store, not an upload), and the window is read as
a prefix view — the cache stops round-tripping entirely.
Tests
tests/unit/test_apple_gpu_resident_kv.py(8): append/window contents,prefill+append, zero-copy aliasing (a later append is visible through an
earlier window — same backing buffer), the resident window feeding a
device-resident bmm (R1) with no host copy + numpy cross-check, single
allocation across 20 appends, overflow raises, footprint,
prefix_viewbounds.Verification (local, Apple Silicon)
torch-import error)The architecture — complete
The whole GPU-resident plan from the scoping doc is now implemented. The one
remaining documented follow-on is a device-resident block-paged cache
(non-contiguous block-table gather on-device) for multi-sequence serving.
🤖 Generated with Claude Code