Add Ollama integration: cross-process KV-cache reuse via the Mooncake Store #5
Loading…
Reference in New Issue
No description provided.
Delete Branch "zbtrs2/Mooncake:feature/ollama-kvcache-bus"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Summary
This adds an integration under
mooncake-integration/ollama/that lets multiple Ollama / llama.cpp processes share one long-prefix KV cache through the Mooncake Store. The first process to prefill a long context publishes its KV to the store; other processes - in a different process, on a different GPU, or on a different node - restore that KV instead of re-prefilling it. The integration is a sidecar plus a thin store proxy and two small, env-gated Ollama patches, so an unconfigured Ollama behaves exactly as before. All changes are confined tomooncake-integration/ollama/(76 files); no existing Mooncake files are modified.Motivation
Ollama is a common entry point for local models and multi-agent tooling, but it has no KV-cache sharing of its own: prefix reuse is confined to a single process's RAM and dies with the request. Mooncake already pools and shares KV for vLLM / SGLang / TRT-LLM. This integration extends that reuse to the local/edge multi-agent case (multiple coding agents sharing one repo context) and to Ollama, with a design that needs almost no change to Ollama itself.
What's included
All paths are relative to
mooncake-integration/ollama/.bridge/- Go sidecar. Serves a gRPCKVCacheBus(Lookup/Prepare/Commit) and an HTTP/JSON gateway with a Prometheus/metricsendpoint. Pure Go (bridge/go.mod, moduleollama-mooncake-bridge, Go 1.23).bridge/internal/cachekey/- GGUF metadata parser and chained per-block content hashing.bridge/internal/prefixindex/- cross-process radix tree of KV prefixes.bridge/internal/arbiter/- restore-vs-recompute cost arbiter with online learning.bridge/internal/orchestrator/- 3-stage Lookup -> Prepare(restore) -> Commit(save).bridge/internal/store/- pluggable backend: Mooncake via proxy or local file.bridge/internal/llamabridge/- Stage-1 llama.cpp/slotsHTTP client.bridge/internal/seqstate/- Stage-2 cgo binding forllama_state_seq_*_ext/ON_DEVICE, behind thecgo_stage2build tag.store-proxy/- Python gRPC wrapper that owns a single warmMooncakeDistributedStorehandle. KV blobs are passed as file paths.ollama-patches/- two additive, env-gated diffs againstollama/ollamaplus notes.benchmarks/,deploy/,scripts/,docs/- benchmark matrix, compose/Prometheus/Grafana, lifecycle scripts, REPORT.md + result JSONs.Design
Three independently deployable processes: agents/patched Ollama -> bridge (Go sidecar: key + radix index + arbiter, 3-stage Lookup/Prepare/Commit) -> store-proxy (one warm store client, striped batch RDMA) -> Mooncake Store.
h_i = H(h_{i-1} || block_i). KV is never mixed across models/tokenizers; favors a miss over a mis-hit.B* = (KV bytes/token) x (prefill tokens/s); learns effective restore bandwidth and prefill rate per model online (EWMA), restores only when cheaper, else recomputes.batch_put_from/batch_get_intoover pre-registered pinned staging buffers. Slot files on/dev/shm./slots; Stage 2 is a cgo path keeping KV in device buffers (targets GPUDirect RDMA).Results
Quoted from the submitted technical report and result artifacts; measured on 2x NVIDIA A100-80GB SXM (NVSwitch), 200 Gb/s InfiniBand, CUDA 12.x, CUDA llama.cpp, 7B-class GGUF models, and q8_0 KV unless noted designed. A 1.5B model is used only for the "restore is not worth it" arbiter check.
Scaling, 7B / 30k ctx, fresh namespace per point (measured):
Single uncontended cross-GPU reuser: honest TTFT 5.99 s (cold) -> 1.83 s (restore), -69%. This isolates the restore path from multi-agent contention and shows the upper bound before the single store client / NIC starts serializing concurrent restores.
Cold-start end-to-end demo (6 agents, 30k context): no-sharing prefill 173,653 tokens -> 28,998 tokens with Mooncake, saving 144,655 tokens; total prefill drops by 83% and redundant prefill is eliminated. Mean TTFT improves from 4,279 ms to 2,934 ms, swarm throughput is 1.26x, and 5 of 6 agents hit the shared KV because agent 0 is the cold cache builder.
Warm RDMA store (same 6-agent shape, public prefix already committed): mean TTFT improves from 4.03 s to 2.88 s (-29%) and swarm throughput reaches 1.74x. This is the steady-state number for repeated work over the same repository context; the cold-start number above is the first-round cost that includes cache creation.
Stage-2 KV-state export microbench: 7B/16k, 875 MiB KV,
/slotsfile path 417 ms vs in-processllama_state_seq_get/set_data_ext91 ms; 1.5B/8k, 219 MiB KV, 105 ms vs 19 ms. Both sizes pass round-trip correctness by continuing decode and matching the next-token argmax. On-device GPUDirect transfer through the Transfer Engine is designed, not yet wired (needs C++ headers absent from the Python wheel).How to test
Go unit tests run anywhere (no GPU/RDMA) and pass:
Full stack + benchmarks need the A100 + InfiniBand setup above, CUDA llama.cpp, the model artifacts, and a reachable Mooncake Store; scripts keep state under
$OMB_STATE:bash scripts/setup_*.sh,bash scripts/demo.sh 7b 6,python benchmarks/run_matrix.py ...,bash scripts/stack_down.sh.Notes / limitations
0002(env-gated--slot-save-path) is a clean standalone upstream candidate;0001adds a self-containedoptions.mooncake.*client without modifying upstream symbols.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.Step 1:
From your project repository, check out a new branch and test the changes.Step 2:
Merge the changes and update on Gitea.