Add Ollama integration: cross-process KV-cache reuse via the Mooncake Store #5

Open

zbtrs2 wants to merge 9 commits from zbtrs2/Mooncake:feature/ollama-kvcache-bus into main

Author	SHA1	Message	Date
zbtrs2	ca8674f51b	test(store): normalize allocator test headings Auto Label PRs / triage (pull_request) Failing after 1m19s Details Build & Test (Linux) / build (3.10) (pull_request) Has been cancelled Details Build & Test (Linux) / build (3.12) (pull_request) Has been cancelled Details Build & Test (Linux) / build-musa (pull_request) Has been cancelled Details Build & Test (Linux) / test-wheel-ubuntu (3.10, ubuntu-22.04) (pull_request) Has been cancelled Details Build & Test (Linux) / test-wheel-ubuntu (3.10, ubuntu-24.04) (pull_request) Has been cancelled Details Build & Test (Linux) / test-wheel-ubuntu (3.12, ubuntu-22.04) (pull_request) Has been cancelled Details Build & Test (Linux) / test-wheel-ubuntu (3.12, ubuntu-24.04) (pull_request) Has been cancelled Details Build & Test (Linux) / build-flags (3.10) (pull_request) Has been cancelled Details Build & Test (Linux) / build-flags (3.12) (pull_request) Has been cancelled Details Build & Test (Linux) / Build Docker Image (pull_request) Has been cancelled Details Build & Test (Linux) / Spell Check with Typos (pull_request) Has been cancelled Details Build & Test (Linux) / Check code format (pull_request) Has been cancelled Details Build & Test (Linux) / check-paths (pull_request) Has been cancelled Details Build & Test (Linux) / build-wheel-cu13 (pull_request) Has been cancelled Details Build & Test (Linux) / ascend-test (pull_request) Has been cancelled Details Build & Test (Linux) / integration-test (pull_request) Has been cancelled Details Build & Test (Linux) / CI Gate (pull_request) Has been cancelled Details	2026-06-29 22:29:34 +08:00
zbtrs2	fd7bbf720a	ollama-kvcache: add deployment, observability and docs deploy packages the topology as docker-compose (master, store proxy, sidecar, patched Ollama, Prometheus, Grafana) with the RDMA/GPUDirect host passthrough it needs, a Prometheus scrape config and a Grafana dashboard whose primary panel integrates the live count of avoided prefill tokens. docs/REPORT.md is the full write-up: design, what is real versus simulated, the measured store bandwidth, the multi-agent results, the adaptive-arbiter behaviour, the Stage-2 microbenchmark and the honest limitations. The README is the entry point and quickstart, and the figures and result JSONs are the measured data behind the report.	2026-06-29 02:13:48 +08:00
zbtrs2	897a11bd34	ollama-kvcache: add benchmarks and orchestration scripts agent_swarm drives the core scenario: several agents share a long repo context and differ only in a short instruction. It runs the workload with sharing off (baseline) and on, erasing llama's own per-slot cache between runs so the comparison isolates cross-process reuse, and reports honest end-to-end TTFT, recomputed prefill tokens, hit rate and throughput. smoke_e2e is a minimal two-agent cross-GPU reuse check; microbench_store and microbench_striped measure raw and striped store bandwidth; run_matrix sweeps a config and plot renders the figures. scripts/env.sh keeps every toolchain, cache and runtime directory under an out-of-tree state root so the source stays clean and nothing is written under /. The setup_* scripts install a local Go toolchain, a venv with mooncake-transfer-engine, a CUDA build of llama.cpp and the GGUF models; the *_start/stop and demo scripts bring the stack up and down, scoped to this integration's own ports and paths.	2026-06-29 02:13:48 +08:00
zbtrs2	e31c2de9b8	ollama-kvcache: add Ollama integration patches Two small, additive patches that make upstream Ollama a participant of the KVCache Bus, verified against ollama/ollama main at 1abd56b. 0001 adds a self-contained HTTP client for the sidecar's prepare/commit endpoints plus parsing of the options.mooncake.* request extension (enable, namespace, read, write, block_size, transport, replica_num, soft_pin), without touching any existing symbol. 0002 injects --slot-save-path into the bundled llama-server command line when OLLAMA_MOONCAKE_SLOT_SAVE_PATH is set, exposing the /slots save\|restore endpoints the sidecar drives. Both are gated, so an unconfigured Ollama behaves exactly as before. The README documents how to apply them and the one call-site hook left as a reference, and notes 0002 as a clean upstream-PR candidate aligned with Ollama issue 14872.	2026-06-29 01:51:49 +08:00
zbtrs2	0b18e65193	ollama-kvcache: add Stage-2 on-device KV state path The Stage-1 path serialises a sequence's KV through host memory and a file (llama.cpp /slots save), which is several times slower than a raw in-process copy because of the double copy described in llama.cpp issue 8915. Stage 2 targets the high-performance route. seqstate is a cgo binding (build tag cgo_stage2) for llama_state_seq_get/set_data_ext. With LLAMA_STATE_SEQ_FLAGS_ON_DEVICE the KV tensors stay in device buffers, ready to be registered with the Mooncake Transfer Engine for GPUDirect RDMA with no host copy. It is gated behind a build tag because it links libllama and only applies to the embedded deployment; the default HTTP path needs none of it. omb_kvbench is an in-process microbenchmark against libllama that compares the file-save path, the raw host export and the on-device size, and verifies that a sequence's KV round-trips correctly (export seq 0, import into seq 1, identical next-token argmax). It quantifies the file-path overhead the on-device route removes.	2026-06-29 01:51:33 +08:00
zbtrs2	56efccf73c	ollama-kvcache: add Mooncake store proxy A long-lived gRPC service that owns one warm MooncakeDistributedStore handle and exposes it to the Go sidecar, which cannot link the Python bindings directly. Using the official store client keeps the integration on the supported API surface. KV snapshots are multi-GiB, so the proxy moves them by file path: PutFile reads a slot file and stores it, GetFile materialises an object straight into the slot directory, and the blob is copied at most once. Objects are striped into <=64 MiB chunks transferred in parallel with batch_put_from / batch_get_into, because single-object RDMA degrades for large objects while striped batches sustain full bandwidth. A pool of pre-registered (pinned) staging buffers amortises RDMA memory registration across operations. Eviction or lease-expiry races on read are reported as not-found so the caller falls back to recompute rather than failing. A filesystem backend (no master required) backs the local baseline, and a selftest verifies put/get round-trips and existence checks over both backends.	2026-06-29 01:51:08 +08:00
zbtrs2	0b208be4e6	ollama-kvcache: add sidecar orchestrator, llama.cpp client and server Wire the pieces into the bridged daemon. orchestrator implements the three-stage reuse flow. Lookup builds the per-block keys and finds the longest stored prefix with a single batched existence query (correct across nodes). Prepare additionally restores the matched KV into a target llama.cpp slot when the arbiter approves, so the server only prefills the uncached tail. Commit saves the slot KV and stores it under the block key for the largest block-aligned prefix, with skip-if-exists giving single-writer dedup when concurrent agents share a prefix. It feeds the arbiter the full restore wall time and the observed prefill rate so estimates self-calibrate, and tracks per-model KV bytes/token for transfer sizing. llamabridge drives the stock llama.cpp server over HTTP: /tokenize, /completion (reading timings.prompt_n to count avoided prefill), /slots save\|restore\|erase and /props. No llama.cpp source change is needed for this path. metrics exposes Prometheus instrumentation; the headline series is the cumulative count of prompt tokens not re-prefilled thanks to reuse. server adapts the orchestrator onto the gRPC KVCacheBus and an HTTP/JSON gateway, and cmd/bridged is the daemon entrypoint serving gRPC (TCP and unix socket), the JSON gateway and /metrics.	2026-06-29 01:50:44 +08:00
zbtrs2	4927b5a694	ollama-kvcache: add gRPC contracts and store backends Define the two service contracts and the object-store abstraction the orchestrator builds on. bridge.proto is the agent-facing KVCacheBus: Lookup (read-only longest prefix match plus arbiter decision), Prepare (Lookup plus restore of the matched KV into a llama.cpp slot) and Commit (save the slot KV back to the store). storeproxy.proto is the sidecar-to-proxy contract; KV blobs are addressed by file path so a multi-GiB snapshot is copied at most once and never streams through the Go process. The store package exposes a small Backend interface with two implementations: a Mooncake backend that forwards file-path Put/Get and batched existence checks to the Python store proxy over gRPC, and a local filesystem backend used as the cross-process baseline and a dependency-free fallback. Generated stubs are checked in so the module builds without protoc.	2026-06-29 01:50:25 +08:00
zbtrs2	2febfc17e6	ollama-kvcache: add cache-key, prefix index and cost arbiter Introduce the algorithmic core of the Ollama KVCache Bus, an integration that lets independent Ollama/llama.cpp processes share prompt KV through the Mooncake Store. cachekey derives a content-addressed key from a model fingerprint (model digest, tokenizer/RoPE hashes parsed from GGUF metadata, KV dtype/layout, context length, block size) and a forward-chained per-block hash of the prompt tokens, where h_i = H(h_{i-1} \|\| block_i). Chaining makes longest-prefix matching exact: two prompts agree on h_i iff they share every token of blocks 0..i. KV produced under different model parameters lands in a disjoint key space and is never reused. prefixindex maintains an in-memory radix tree over those chained hashes, the cross-process analogue of a RadixAttention tree whose nodes reference remote KV snapshots. It answers longest-prefix lookups in O(blocks) and tracks per-prefix fan-out (hotness) and recency for replication and eviction decisions. arbiter decides restore-vs-recompute per request. A cached prefix only helps when the store can deliver its KV faster than the GPU regenerates it; the arbiter estimates both from rates it learns online (EWMA, seeded on first observation) per model, and restores only when it is cheaper. This keeps KV sharing loss-free across hardware regimes. Includes unit tests for hash determinism, prefix consistency, model separation, longest-prefix matching, eviction ordering and the arbiter's adapt-to-recompute behaviour.	2026-06-29 01:50:09 +08:00