[Store] Add Mooncake KV Event Publisher for Dynamo-compatible KV events

mooncake-integration/store/store_py.cpp

@@ -1786,6 +1786,36 @@ PYBIND11_MODULE(store, m) {
         .value("GENERAL", ObjectDataType::GENERAL)
         .export_values();
 
+    // KV event metadata types attached to writes via ReplicateConfig.
+    py::class_<KvBlockComponentSpec>(m, "KvBlockComponentSpec")
+        .def(py::init<>())
+        .def_readwrite("object_key", &KvBlockComponentSpec::object_key)
+        .def_readwrite("component_role", &KvBlockComponentSpec::component_role)
+        .def_readwrite("component_index", &KvBlockComponentSpec::component_index);
+
+    py::class_<KvBlockEventMetadata>(m, "KvBlockEventMetadata")
+        .def(py::init<>())
+        .def_readwrite("schema_version", &KvBlockEventMetadata::schema_version)
+        .def_readwrite("group_id", &KvBlockEventMetadata::group_id)
+        .def_readwrite("block_hash", &KvBlockEventMetadata::block_hash)
+        .def_readwrite("parent_block_hash",
+                       &KvBlockEventMetadata::parent_block_hash)
+        .def_readwrite("token_ids", &KvBlockEventMetadata::token_ids)
+        .def_readwrite("block_size", &KvBlockEventMetadata::block_size)
+        .def_readwrite("dp_rank", &KvBlockEventMetadata::dp_rank)
+        .def_readwrite("model_name", &KvBlockEventMetadata::model_name)
+        .def_readwrite("lora_name", &KvBlockEventMetadata::lora_name)
+        .def_readwrite("additional_salt",
+                       &KvBlockEventMetadata::additional_salt)
+        .def_readwrite("expected_object_count",
+                       &KvBlockEventMetadata::expected_object_count)
+        .def_readwrite("expected_components",
+                       &KvBlockEventMetadata::expected_components)
+        .def_readwrite("emit_stored_event",
+                       &KvBlockEventMetadata::emit_stored_event)
+        .def_readwrite("emit_removed_event",
+                       &KvBlockEventMetadata::emit_removed_event);
+
     // Define the ReplicateConfig class
     py::class_<ReplicateConfig>(m, "ReplicateConfig")
         .def(py::init<>())
@@ -1802,6 +1832,8 @@ PYBIND11_MODULE(store, m) {
                        &ReplicateConfig::prefer_alloc_in_same_node)
         .def_readwrite("data_type", &ReplicateConfig::data_type)
         .def_readwrite("group_ids", &ReplicateConfig::group_ids)
+        .def_readwrite("kv_event_metadata",
+                       &ReplicateConfig::kv_event_metadata)
         .def("__str__", [](const ReplicateConfig &config) {
             std::ostringstream oss;
             oss << config;

mooncake-store/include/kv_event.h

@@ -0,0 +1,115 @@
+#pragma once
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace mooncake {
+
+/**
+ * @brief Type of a Dynamo-compatible KV cache event published by Mooncake.
+ */
+enum class KvEventType {
+    kBlockStored,
+    kBlockRemoved,
+};
+
+/**
+ * @brief A single Dynamo-compatible KV cache event.
+ *
+ * Field semantics follow the Dynamo ZMQ relay "map event" wire format plus the
+ * Mooncake extension fields (source / tenant_id / group_id / worker_id /
+ * event_id). The encoder in kv_event.cpp serializes this into the msgpack map
+ * that the Dynamo Mooncake adapter consumes.
+ *
+ * BlockStored uses all fields; BlockRemoved only needs the identity fields plus
+ * block_hashes and medium (token_ids/block_size/parent_block_hash are ignored).
+ */
+struct KvEvent {
+    KvEventType type{KvEventType::kBlockStored};
+
+    // Identity / Mooncake extension fields.
+    std::string event_id;
+    std::string source{"mooncake"};
+    std::string tenant_id;
+    std::string group_id;
+    std::string worker_id;
+
+    // Dynamo core fields.
+    std::vector<uint64_t> block_hashes;
+    std::string medium{"EXTERNAL"};
+
+    // BlockStored-only fields.
+    std::optional<uint64_t> parent_block_hash;
+    std::vector<uint32_t> token_ids;
+    uint32_t block_size{0};
+    std::string lora_name;
+    std::string model_name;
+    std::optional<uint32_t> dp_rank;
+
+    bool operator==(const KvEvent& other) const {
+        return type == other.type && event_id == other.event_id &&
+               source == other.source && tenant_id == other.tenant_id &&
+               group_id == other.group_id && worker_id == other.worker_id &&
+               block_hashes == other.block_hashes && medium == other.medium &&
+               parent_block_hash == other.parent_block_hash &&
+               token_ids == other.token_ids && block_size == other.block_size &&
+               lora_name == other.lora_name && model_name == other.model_name &&
+               dp_rank == other.dp_rank;
+    }
+};
+
+/**
+ * @brief Abstract sink for Dynamo-compatible KV events.
+ *
+ * The master calls Publish() when a logical KV block transitions between
+ * complete/incomplete. Implementations decide the transport (e.g. a single ZMQ
+ * PUB stream). Tests use an in-memory mock. A null publisher means events are
+ * disabled and Mooncake behavior is unchanged.
+ */
+class KvEventPublisher {
+   public:
+    virtual ~KvEventPublisher() = default;
+    virtual void Publish(const KvEvent& event) = 0;
+};
+
+// --- Encoding / decoding (msgpack) -----------------------------------------
+//
+// These produce / consume the Dynamo ZMQ relay wire format. Returning byte
+// strings keeps msgpack out of widely-included headers.
+
+/**
+ * @brief Encode one event into the Dynamo "map event" msgpack representation.
+ */
+std::string EncodeKvEventMap(const KvEvent& event);
+
+/**
+ * @brief Encode a batch of events into the ZMQ relay payload:
+ * [timestamp (f64), [events], dp_rank (i32, optional)].
+ */
+std::string EncodeKvEventBatchPayload(const std::vector<KvEvent>& events,
+                                      double timestamp,
+                                      std::optional<int32_t> dp_rank);
+
+/**
+ * @brief Decode one msgpack "map event". Returns std::nullopt on malformed
+ * input or unsupported event type. Intended for tests / adapters.
+ */
+std::optional<KvEvent> DecodeKvEventMap(const std::string& bytes);
+
+/**
+ * @brief Result of decoding a ZMQ relay batch payload.
+ */
+struct KvEventBatch {
+    double timestamp{0.0};
+    std::vector<KvEvent> events;
+    std::optional<int32_t> dp_rank;
+};
+
+/**
+ * @brief Decode a ZMQ relay payload produced by EncodeKvEventBatchPayload.
+ */
+std::optional<KvEventBatch> DecodeKvEventBatchPayload(const std::string& bytes);
+
+}  // namespace mooncake

mooncake-store/include/kv_event_zmq_publisher.h

@@ -0,0 +1,72 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+
+#include "kv_event.h"
+
+namespace mooncake {
+
+/**
+ * @brief Configuration for ZmqKvEventPublisher.
+ */
+struct ZmqKvEventPublisherConfig {
+    // ZMQ endpoint to publish on, e.g. "tcp://0.0.0.0:5557". Use a "*" port
+    // (e.g. "tcp://127.0.0.1:*") to let ZMQ choose a free port; the chosen
+    // endpoint can be read back via ZmqKvEventPublisher::endpoint().
+    std::string endpoint{"tcp://0.0.0.0:5557"};
+    // ZMQ topic prefix sent as frame 1. Dynamo's listener matches this against
+    // its zmq_topic subscription filter; default is empty (matches all).
+    std::string topic{};
+    // When true (default) the PUB socket binds the endpoint (Mooncake master is
+    // the stable server). When false it connects instead (for relays/proxies).
+    bool bind{true};
+    // Linger in milliseconds applied on close so a clean shutdown can flush
+    // buffered events. -1 keeps the ZMQ default (block forever); 0 drops.
+    int linger_ms{0};
+    // Send high-water-mark (max queued outbound messages); 0 means unlimited.
+    int send_hwm{0};
+};
+
+/**
+ * @brief KvEventPublisher backed by a single ZMQ PUB socket.
+ *
+ * Emits the SGLang/vLLM-compatible 3-frame message Dynamo's ZMQ relay expects:
+ *   frame 1: topic
+ *   frame 2: 8-byte big-endian monotonic sequence number
+ *   frame 3: msgpack payload [timestamp, [event], dp_rank]
+ *
+ * Because Mooncake's master is the single centralized owner of KV state, all
+ * workers' events flow through this one PUB stream; the physical owner of each
+ * block is carried in the per-event worker_id field (see kv_event.h). Each
+ * Publish() call sends a single-event batch.
+ *
+ * Thread-safety: Publish() is serialized with an internal mutex, so the
+ * publisher may be shared across master threads.
+ */
+class ZmqKvEventPublisher : public KvEventPublisher {
+   public:
+    explicit ZmqKvEventPublisher(const ZmqKvEventPublisherConfig& config);
+    ~ZmqKvEventPublisher() override;
+
+    ZmqKvEventPublisher(const ZmqKvEventPublisher&) = delete;
+    ZmqKvEventPublisher& operator=(const ZmqKvEventPublisher&) = delete;
+
+    void Publish(const KvEvent& event) override;
+
+    // The actual bound/connected endpoint. When the config used a wildcard
+    // port (e.g. "tcp://127.0.0.1:*"), this returns the concrete endpoint ZMQ
+    // selected (e.g. "tcp://127.0.0.1:54321").
+    const std::string& endpoint() const;
+
+    // Number of messages successfully handed to ZMQ so far (for diagnostics).
+    uint64_t published_count() const;
+
+   private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace mooncake