[Store] feat: add workload-aware KV cache hints by Lin-z-w · Pull Request #2830 · kvcache-ai/Mooncake

Lin-z-w · 2026-07-10T07:18:53Z

Description

Implements Stage 1 of RFC #2819: Workload-Aware KV Cache Hints. This PR does not close the RFC because eviction policy, observability, and external framework integrations remain follow-up work.

This change:

adds a nested WorkloadHints API with optional session_id and retention_priority fields while leaving all existing ReplicateConfig fields in place;
propagates workload hints through C++/Python write APIs, batch sharding, tensor write paths, and client-to-master RPCs;
stores hints as immutable object metadata without changing placement, lease, pinning, or eviction behavior;
applies one set of hints to every object in a batch and preserves existing group_ids semantics;
uses insert-only metadata semantics for duplicate PUT operations and preserves the original hints when an existing object is upserted;
persists hints in HA snapshots and restores API defaults from pre-hints snapshot formats; and
keeps existing Python call patterns compatible when workload hints are unset.

retention_priority is stored as provided. Callers or integration layers are responsible for normalization. Same-version client/master RPC propagation is covered here; rolling mixed-version RPC compatibility is not introduced in this stage.

Module

Type of Change

How Has This Been Tested?

Coverage includes default and non-default hints, single and batch writes, coexistence with group_ids, duplicate PUT, same-size and size-changing upserts, real client/master RPC propagation, current and legacy snapshot restoration, catalog-backed snapshot acceptance, and Python binding round trips.

Test commands:

cmake --build build --target \
  master_service_test client_integration_test pybind_client_test \
  snapshot_child_process_test catalog_backed_snapshot_provider_test store -j 8

ctest --test-dir build \
  -R '^(master_service_test|client_integration_test|pybind_client_test|snapshot_child_process_test)$' \
  --output-on-failure -j 2

build/mooncake-store/tests/catalog_backed_snapshot_provider_test

PYTHONPATH=build/mooncake-integration python3 - <<'PY'
from store import ReplicateConfig, WorkloadHints

hints = WorkloadHints()
hints.session_id = "session-a"
hints.retention_priority = -7
config = ReplicateConfig()
config.workload_hints = hints
assert config.workload_hints.session_id == "session-a"
assert config.workload_hints.retention_priority == -7
PY

pre-commit run --files $(git diff --name-only upstream/main)
git diff --check upstream/main...HEAD

cd docs
make html SPHINXBUILD=.venv/bin/sphinx-build

Test results:

Unit tests pass
Integration tests pass (if applicable)
Manual testing done (describe below)
The four requested CTest targets passed: 4/4.
Catalog-backed snapshot provider tests passed: 9/9.
Python import and binding tests passed: 3/3, plus a native-extension field round trip.
Targeted pre-commit hooks and git diff --check passed.
The Sphinx HTML build completed successfully. It emitted seven warnings because external intersphinx inventories were unreachable in the restricted network environment.

Checklist

I have performed a self-review of my own code
I have formatted my code using ./scripts/code_format.sh
I have run pre-commit run --all-files and all hooks pass
I have updated the documentation (if applicable)
I have added tests to prove my changes are effective
For changes >500 LOC: I have filed an RFC issue (#2819)

AI Assistance Disclosure

No AI tools were used
AI tools were used (specify below)

OpenAI Codex assisted with implementation, test development, and compatibility review. I reviewed and edited the generated code before submission.

gemini-code-assist

Code Review

This pull request introduces WorkloadHints (consisting of session_id and retention_priority) to associate workload-level metadata with stored objects, integrating it across ReplicateConfig, ObjectMetadata, MasterService, and Python bindings. It also updates MsgPack serialization/deserialization to support these hints while maintaining backward compatibility with older snapshot formats. Feedback on these changes suggests removing noexcept from the WorkloadHints stream insertion operator to prevent unsafe terminations, optimizing ReplicateConfig::IsDefault() by comparing against a static const instance, and improving forward-compatibility during metadata deserialization by verifying the element type before parsing it as workload hints.

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.

gemini-code-assist · 2026-07-10T07:27:25Z

+    friend std::ostream& operator<<(std::ostream& os,
+                                    const WorkloadHints& hints) noexcept {


Marking operator<< as noexcept is unsafe because stream insertion operations (such as writing to os or copying/formatting hints.session_id) can throw exceptions (e.g., if the stream has its exception mask configured to throw on failure, or due to allocation failures). If an exception is thrown inside a noexcept function, std::terminate will be called immediately. It is standard practice to omit noexcept from stream insertion operators.

Suggested change

friend std::ostream& operator<<(std::ostream& os,

const WorkloadHints& hints) noexcept {

friend std::ostream& operator<<(std::ostream& os,

const WorkloadHints& hints) {

gemini-code-assist · 2026-07-10T07:27:25Z

+
+    bool operator==(const ReplicateConfig&) const = default;
+
+    bool IsDefault() const { return *this == ReplicateConfig{}; }


To avoid constructing a temporary ReplicateConfig object (which may involve initializing strings, vectors, and optionals) on every call to IsDefault(), you can compare against a static const instance.

bool IsDefault() const { static const ReplicateConfig kDefault{}; return *this == kDefault; }

gemini-code-assist · 2026-07-10T07:27:25Z

+    if (index < total_elements) {
+        const auto& hints_object = array[index++];
+        if (hints_object.type != msgpack::type::ARRAY ||
+            hints_object.via.array.size != 2) {


To ensure robust forward-compatibility, we should check if the next element is actually an array before consuming it as workload_hints. If a future version of the snapshot format omits workload_hints but appends other optional fields of different types, unconditionally advancing index and expecting an array will cause deserialization to fail. Checking the type first allows the deserializer to safely skip workload_hints and ignore unknown trailing fields.

if (index < total_elements && array[index].type == msgpack::type::ARRAY) { const auto& hints_object = array[index++]; if (hints_object.via.array.size != 2) {

codecov-commenter · 2026-07-10T08:53:20Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 96.67897% with 9 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-integration/store/store_py.cpp	0.00%	5 Missing ⚠️
mooncake-store/tests/master_service_test.cpp	98.58%	2 Missing ⚠️
mooncake-integration/store/store_py_internal.h	0.00%	1 Missing ⚠️
mooncake-store/src/master_service.cpp	96.55%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[Store] feat: add workload-aware KV cache hints

c59aa2a

Lin-z-w requested review from 00fish0, Libotry, ShangmingCai, XucSh, YiXR, stmatengss and ykwd as code owners July 10, 2026 07:18

gemini-code-assist Bot reviewed Jul 10, 2026

View reviewed changes

github-actions Bot added run-ci Store Installation Integration Tests labels Jul 10, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Store] feat: add workload-aware KV cache hints#2830

[Store] feat: add workload-aware KV cache hints#2830
Lin-z-w wants to merge 1 commit into
kvcache-ai:mainfrom
Lin-z-w:feat/workload-aware-kv-hints

Lin-z-w commented Jul 10, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Uh oh!

codecov-commenter commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		friend std::ostream& operator<<(std::ostream& os,
		const WorkloadHints& hints) noexcept {


		bool operator==(const ReplicateConfig&) const = default;

		bool IsDefault() const { return *this == ReplicateConfig{}; }

Uh oh!

Conversation

Lin-z-w commented Jul 10, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Checklist

AI Assistance Disclosure

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jul 10, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Lin-z-w commented Jul 10, 2026 •

edited by github-actions Bot

Loading