[Store] Add Ollama connector for PD-disaggregated KV cache sharing#2800
[Store] Add Ollama connector for PD-disaggregated KV cache sharing#2800Lfan-ke wants to merge 1 commit into
Conversation
Lfan-ke
requested review from
ShangmingCai,
XucSh,
YiXR,
stmatengss and
ykwd
as code owners
github-actions
Bot
added
documentation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the Mooncake x Ollama integration, adding an Ollama-side connector to offload and reuse KV cache through Mooncake Store, along with configuration, prefix-keying logic, an end-to-end example, and documentation. Feedback on the implementation suggests keeping c.store non-nil in Close() to prevent potential nil pointer dereferences in post-close calls, optimizing the Pages function by reusing a single SHA-256 hasher to reduce heap allocations, removing the unused chain helper, and either removing or utilizing the unused ErrBufferCount error.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| func (c *Connector) Close() { | ||
| if c.store != nil { | ||
| c.store.Close() | ||
| c.store = nil | ||
| } | ||
| } |
There was a problem hiding this comment.
Setting c.store = nil in Close() introduces a risk of nil pointer dereference panics if any of the connector's other methods (such as RegisterKVBuffer, LoadCachedPrefix, or StoreComputedPrefix) are called after the connector is closed.
If we keep c.store non-nil, any post-close calls will safely delegate to c.store's methods. Since mooncakestore.Store.Close() sets its internal handle to nil, and all Store methods check if s.handle == nil { return ErrStoreNil }, they will gracefully return ErrStoreNil instead of panicking. This leverages the existing robustness of the underlying store client without requiring redundant nil checks in every method of Connector.
func (c *Connector) Close() {
if c.store != nil {
c.store.Close()
}
}| func (c Config) Pages(tokens []int32) ([]Page, error) { | ||
| if c.PageSize <= 0 { | ||
| return nil, ErrInvalidPageSize | ||
| } | ||
|
|
||
| n := NumFullPages(len(tokens), c.PageSize) | ||
| pages := make([]Page, 0, n) | ||
| prefixStr := c.keyPrefix() | ||
|
|
||
| h := sha256.Sum256([]byte(hashSeedDomain + c.Model)) | ||
| for i := 0; i < n; i++ { | ||
| start := i * c.PageSize | ||
| end := start + c.PageSize | ||
| h = chain(h, tokens[start:end]) | ||
| key := fmt.Sprintf("%s/%s/%s", prefixStr, c.Model, hex.EncodeToString(h[:keyHashBytes])) | ||
| pages = append(pages, Page{Index: i, StartToken: start, EndToken: end, Key: key}) | ||
| } | ||
| return pages, nil | ||
| } |
There was a problem hiding this comment.
In high-throughput LLM serving, prompt processing latency (TTFT) is highly critical. Currently, Pages calls chain for every single page, which allocates a new SHA-256 hasher (sha256.New()) on the heap each time. For long prompts with many pages, this introduces significant heap allocation overhead and garbage collection pressure.
We can optimize this by creating a single hash.Hash instance at the start of Pages and calling Reset() on it for each page. Additionally, pre-allocating the pages slice with its exact size n and assigning elements directly by index (instead of using append) avoids slice growth overhead.
func (c Config) Pages(tokens []int32) ([]Page, error) {
if c.PageSize <= 0 {
return nil, ErrInvalidPageSize
}
n := NumFullPages(len(tokens), c.PageSize)
pages := make([]Page, n)
prefixStr := c.keyPrefix()
hasher := sha256.New()
hasher.Write([]byte(hashSeedDomain + c.Model))
h := hasher.Sum(nil)
var b [4]byte
for i := 0; i < n; i++ {
hasher.Reset()
hasher.Write(h)
start := i * c.PageSize
end := start + c.PageSize
for _, t := range tokens[start:end] {
binary.BigEndian.PutUint32(b[:], uint32(t))
hasher.Write(b[:])
}
h = hasher.Sum(nil)
key := fmt.Sprintf("%s/%s/%s", prefixStr, c.Model, hex.EncodeToString(h[:keyHashBytes]))
pages[i] = Page{Index: i, StartToken: start, EndToken: end, Key: key}
}
return pages, nil
}| // chain folds one page of tokens into the running digest. Tokens are encoded as | ||
| // fixed-width big-endian uint32 (llama.cpp token ids are int32) so the encoding | ||
| // is unambiguous and stable across platforms. | ||
| func chain(prev [32]byte, tokens []int32) [32]byte { | ||
| hasher := sha256.New() | ||
| hasher.Write(prev[:]) | ||
| var b [4]byte | ||
| for _, t := range tokens { | ||
| binary.BigEndian.PutUint32(b[:], uint32(t)) | ||
| hasher.Write(b[:]) | ||
| } | ||
| var out [32]byte | ||
| copy(out[:], hasher.Sum(nil)) | ||
| return out | ||
| } |
There was a problem hiding this comment.
Since the hashing logic has been inlined and optimized directly inside Pages to reuse a single hasher, the unexported helper function chain is no longer needed and can be safely removed.
| // ErrBufferCount is returned when the supplied per-page buffer slice is too short. | ||
| ErrBufferCount = errors.New("ollama: fewer page buffers than pages") |
There was a problem hiding this comment.
The exported error ErrBufferCount is defined here but is never used anywhere in the package. If silent truncation of pages in keysFor is the intended behavior, this unused error should be removed to keep the codebase clean. Alternatively, if you want to prevent silent truncation when the caller provides fewer buffers than there are full pages in tokens, you should use this error to validate the buffer slice lengths in LoadCachedPrefix and StoreComputedPrefix.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
To make this connector easy to try during review, I added a small self-contained demo that exercises it end to end on a single machine with no RDMA.
It starts a real Mooncake Store ( Reference run (realistic GQA layout, 2048-token prompt with a 1792-token shared prefix):
Scope, kept explicit in |
|
Thanks for the contribution! Could you share how you envision this being used in practice? Since upstream Ollama doesn't currently support PD-disaggregated inference, I am curious about the intended integration path. |
Ollama keeps KV cache only in its own process today, so it cannot reuse a prefix computed elsewhere or take part in a prefill/decode-disaggregated deployment. This adds an Ollama-side connector on top of the Go store bindings, following the same hash-based prefix-caching pattern as the vLLM MooncakeStoreConnector and SGLang HiCache L3 backend. - mooncake-store/go/ollama/prefix: pure-Go, content-addressed page keying (SHA-256 hash chain over fixed-size token pages, seeded with the model name). Shared token prefixes map to identical keys; any earlier divergence changes every later key, so a worker never loads KV that mismatches its own tokens. Fully unit-tested without a running cluster. - mooncake-store/go/ollama: Connector with prefill/decode/mixed roles. LoadCachedPrefix loads the longest cached prefix via BatchExists + BatchGetInto; StoreComputedPrefix publishes pages via BatchPutFrom, skipping pages a peer already stored. Reuses the existing zero-copy buffer path. - examples/ollama: end-to-end demo over the TCP transport that publishes a prompt's KV, reuses it for a request sharing the prefix, and verifies the bytes round-trip. - docs: Mooncake x Ollama integration guide. Wiring the page buffers to llama.cpp's KV tensors and scheduler-level PD routing live in the Ollama runtime and are tracked as follow-ups; this PR is the store-transfer and keying half of the integration. Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>
Lfan-ke
force-pushed
the
feat/ollama-pd-integration
branch
from
c04abf3 to
895a197
Compare
3 participants
Description
Ollama keeps its KV cache only in-process today, so it cannot reuse a prefix computed elsewhere or take part in a prefill/decode-disaggregated (PD) deployment. This adds an Ollama-side connector on top of the Go store bindings, following the same hash-based prefix-caching pattern as the vLLM
MooncakeStoreConnectorand the SGLang HiCache L3 backend: a prefill worker publishes the KV it computes, and a decode worker (or any peer) reuses the longest cached prefix instead of recomputing it.New code, all under
mooncake-store/go:ollama/prefix— pure-Go, content-addressed page keying: a SHA-256 hash chain over fixed-size token pages, seeded with the model name. Shared token prefixes map to identical keys; any earlier divergence changes every later key, so a worker never loads KV that mismatches its own tokens. No CGo/store dependency, so it is fully unit-tested without a running cluster.ollama—Connectorwith prefill / decode / mixed roles.LoadCachedPrefixloads the longest cached prefix viaBatchExists+BatchGetInto;StoreComputedPrefixpublishes pages viaBatchPutFrom, skipping pages a peer already stored. Reuses the existing zero-copy registered-buffer path.examples/ollama— an end-to-end demo over the TCP transport that publishes a prompt's KV, reuses it for a request sharing the prefix, and verifies the bytes round-trip.Scope: this is the store-transfer and keying half of the integration. Wiring the page buffers to llama.cpp's KV tensors and scheduler-level PD routing live in the Ollama runtime and are tracked as follow-ups.
Module
mooncake-store)Type of Change
How Has This Been Tested?
Test commands:
Test results:
ollama/prefix(keying / hash-chain / longest-present-run) is green.ollamaconnector links the CGomooncakestorebindings, so its build/integration run needs the native store (store_c.h) and a runningmooncake_master; verified where that is available, not in a bare Go environment.examples/ollamapublishes a prompt's KV over TCP, reloads it for a prefix-sharing request, and asserts the bytes round-trip.Checklist
gofmt-clean (scripts/code_format.shtargets C/C++ only; no C/C++ changed here)pre-commit run --all-files— the tool was not available in my environment; the applicable hooks were verified manually (no trailing whitespace, files end with a newline, no large files added)AI Assistance Disclosure
An AI coding assistant was used to help draft the connector, the keying logic and the documentation. I have reviewed, built and tested the changes and take full responsibility for them.
Notes on scope
This is the store-transfer and page-keying half of the integration. Two pieces are intentionally left as follow-ups because they belong in the Ollama runtime, not here:
llama.cpp's paged KV tensors and the per-page buffers. The connector defines the buffer contract (PageLayout+ one pointer per page).