From e7be7299e0da7b443d98cd9a2b62ba3fe5d56087 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 23:01:02 +0800 Subject: [PATCH 01/57] Bump golang.org/x/net in /mooncake-common/k8s-lease (#2742) Bumps [golang.org/x/net](https://github.com/golang/net) from 0.47.0 to 0.55.0. - [Commits](https://github.com/golang/net/compare/v0.47.0...v0.55.0) --- updated-dependencies: - dependency-name: golang.org/x/net dependency-version: 0.55.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- mooncake-common/k8s-lease/go.mod | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mooncake-common/k8s-lease/go.mod b/mooncake-common/k8s-lease/go.mod index 4bfc205264..b3c7f5cdf2 100644 --- a/mooncake-common/k8s-lease/go.mod +++ b/mooncake-common/k8s-lease/go.mod @@ -1,6 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-common/k8s-lease -go 1.24.0 +go 1.25.0 require ( k8s.io/api v0.34.3 @@ -40,11 +40,11 @@ require ( github.com/x448/float16 v0.8.4 // indirect go.yaml.in/yaml/v2 v2.4.3 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/net v0.47.0 // indirect + golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.30.0 // indirect - golang.org/x/sys v0.38.0 // indirect - golang.org/x/term v0.37.0 // indirect - golang.org/x/text v0.31.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/term v0.43.0 // indirect + golang.org/x/text v0.37.0 // indirect golang.org/x/time v0.9.0 // indirect google.golang.org/protobuf v1.36.8 // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect From 1fe59e6a30d00a8fae89735ff7544514502ed57c Mon Sep 17 00:00:00 2001 From: Yihe Liu Date: Sat, 4 Jul 2026 08:04:28 -0700 Subject: [PATCH 02/57] [TE] Add TPU (PJRT) staging support to TENT (#2733) * [TE] Add TPU (PJRT) staging support to TENT Adds Google TPU support to the TENT transfer engine. TPU HBM is not NIC-addressable, so transfers touching TPU memory are staged through host DRAM and chained by the existing ProxyManager pipeline: TPU HBM <-> host DRAM (PJRT device copy) -> host <-> host (RDMA/TCP) Rather than a standalone networked transport, this reuses ProxyManager's chunked double-buffering / staging machinery and adds only: * MTYPE_TPU memory type + "tpu" location/type parsing. * TpuPlatform (derives CpuPlatform): host DRAM + NIC topology are inherited; only the TPU-device-aware paths (copy, getMemoryType, getLocation, per-device MemEntry probe) are overridden and delegated to a device-copy adapter. * A thin TpuTransport: the local HBM<->host staging executor. It only advertises gpu_to_dram / dram_to_gpu (gpu_to_gpu stays false so the engine always stages cross-node traffic through host DRAM) and runs the copy via Platform::copy for LOCAL_SEGMENT_ID requests. It is same-machine-only, like SHM/NVLINK. * A findStagingPolicy case for TPU: local HBM<->host via TpuTransport, host<->host via whichever host transport is present (RDMA or TCP; cloud TPU deployments are typically TCP/multi-NIC). The PJRT device I/O (HBM<->host DMA, pointer classification, device topology) sits behind TpuPjrtShim, which resolves an adapter shared library at runtime via dlopen (C ABI in tpu_pjrt_abi.h). TENT therefore carries no build-time PJRT/XLA dependency. The whole feature is gated behind -DUSE_TPU (OFF by default), so CI and existing builds are unaffected. Includes a mock adapter + unit test that exercise the shim ABI on any Linux host without TPU hardware, and a docs entry under supported protocols. Refs #2662. --- .../getting_started/supported-protocols.md | 34 ++++ mooncake-common/common.cmake | 8 + .../tent/include/tent/common/types.h | 1 + .../tent/include/tent/platform/tpu.h | 66 +++++++ .../tent/include/tent/platform/tpu_pjrt_abi.h | 64 +++++++ .../include/tent/platform/tpu_pjrt_shim.h | 93 ++++++++++ .../tent/include/tent/runtime/platform.h | 5 +- .../tent/transport/tpu/tpu_transport.h | 93 ++++++++++ .../tent/src/CMakeLists.txt | 2 + .../tent/src/platform/CMakeLists.txt | 4 + .../tent/src/platform/tpu/CMakeLists.txt | 7 + .../tent/src/platform/tpu/README.md | 73 ++++++++ .../src/platform/tpu/STAGING_ARCHITECTURE.md | 171 ++++++++++++++++++ .../tent/src/platform/tpu/tpu_pjrt_shim.cpp | 144 +++++++++++++++ .../tent/src/platform/tpu/tpu_platform.cpp | 103 +++++++++++ .../tent/src/runtime/platform.cpp | 4 + .../tent/src/runtime/transfer_engine_impl.cpp | 45 ++++- .../tent/src/runtime/transport_loader.cpp | 9 + .../tent/src/runtime/transport_selector.cpp | 49 +++-- .../tent/src/transport/CMakeLists.txt | 4 +- .../tent/src/transport/tpu/CMakeLists.txt | 5 + .../tent/src/transport/tpu/tpu_transport.cpp | 165 +++++++++++++++++ .../tent/tests/CMakeLists.txt | 23 +++ .../tent/tests/tpu/mock_tpu_pjrt_adapter.cpp | 90 +++++++++ .../tent/tests/tpu/tpu_pjrt_shim_test.cpp | 117 ++++++++++++ 25 files changed, 1361 insertions(+), 18 deletions(-) create mode 100644 mooncake-transfer-engine/tent/include/tent/platform/tpu.h create mode 100644 mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h create mode 100644 mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h create mode 100644 mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt create mode 100644 mooncake-transfer-engine/tent/src/platform/tpu/README.md create mode 100644 mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md create mode 100644 mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp create mode 100644 mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt create mode 100644 mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp create mode 100644 mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp create mode 100644 mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp diff --git a/docs/source/getting_started/supported-protocols.md b/docs/source/getting_started/supported-protocols.md index f98edb2fa8..d74ad0d4a3 100644 --- a/docs/source/getting_started/supported-protocols.md +++ b/docs/source/getting_started/supported-protocols.md @@ -16,6 +16,7 @@ Mooncake Transfer Engine supports multiple communication protocols for data tran | **barex** | RDMA-capable NIC | Bare-metal RDMA extension | ⚠️ Advanced | | **cxl** | CXL-capable hardware | Memory pooling and sharing | ⚠️ Advanced | | **ascend** | Huawei Ascend NPU | Ascend NPU communication | ⚠️ Advanced | +| **tpu** | Google TPU (PJRT) | TPU KV-cache transfer via host-DRAM staging | 🧪 Experimental (TENT) | ## Commonly Used Protocols (Python API) @@ -266,6 +267,39 @@ export MC_FORCE_MNNVL=true - [Heterogeneous Ascend](../design/transfer-engine/heterogeneous_ascend.md) - [Ascend Transport](../design/transfer-engine/ascend_transport.md) +### TPU Transport (tpu) — Experimental + +**Description:** Google TPU support in the TENT runtime. Because TPU HBM is not +directly addressable by the NIC, transfers touching TPU memory are staged +through host DRAM: the HBM ↔ host-DRAM hop is performed by a PJRT device-copy +adapter, and the host ↔ host hop is carried by an existing transport (RDMA/TCP). +The two stages are chained automatically by the TENT staging pipeline +(`ProxyManager`), so no separate networked TPU transport is required. + +**Status:** Experimental. The C++/TENT data path is gated behind `-DUSE_TPU=ON` +(OFF by default). A serving-framework (JAX / PyTorch-XLA) integration layer is +planned as a follow-up. + +**Use When:** +- Disaggregated prefill/decode serving on TPU hosts +- KV-cache transfer between TPU nodes over RDMA/TCP + +**Requirements:** +- Built with `-DUSE_TPU=ON -DUSE_TENT=ON` +- A PJRT device-copy adapter shared library exposing the `mc_tpu_pjrt_*` C ABI + (see `tpu_pjrt_abi.h`). The adapter is resolved at runtime via `dlopen`; its + path defaults to `libmooncake_tpu_pjrt.so` and can be overridden with the + `MC_TPU_PJRT_LIB` environment variable. No PJRT/TPU SDK is required at build + time. +- An RDMA (or TCP) transport enabled for the host ↔ host hop. + +**Design notes:** +- TPU memory is reported as a distinct memory type (`tpu:N` locations); the + staging policy routes the local HBM ↔ host copy to the TPU device-copy + transport and the cross-node hop to RDMA/TCP. +- DMA-mapped (pinned) staging buffers for true async device DMA are a planned + performance follow-up. + ## Configuration Examples ### Configuration File (JSON) diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 2426a6ebe1..fd5292be5d 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -91,6 +91,9 @@ option(USE_EFA "option for using AWS EFA transport" OFF) option(USE_UB "option for using UB protocol transport" OFF) option(USE_SUNRISE "option for enabling gpu features for Sunrise GPU with Tang runtime" OFF) +option(USE_TPU + "option for enabling TPU (PJRT) staging support in TENT; the PJRT adapter is loaded at runtime via dlopen, no build-time SDK required" + OFF) if(USE_UB) add_compile_definitions(USE_UB) @@ -203,6 +206,11 @@ if(USE_CUDA) link_directories(/usr/local/cuda/lib /usr/local/cuda/lib64) endif() +if(USE_TPU) + add_compile_definitions(USE_TPU) + message(STATUS "TPU (PJRT) staging support is enabled") +endif() + if(NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "") if(DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "") set(NEUWARE_ROOT diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index 27a7e98d32..4b7fa80503 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -54,6 +54,7 @@ enum TransportType : int { TCP, AscendDirect, SUNRISE_LINK, + TPU, // Sentinel: must remain the last enumerator. kNumTransportTypes, }; diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu.h new file mode 100644 index 0000000000..1bca92a70d --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu.h @@ -0,0 +1,66 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_PLATFORM_TPU_H_ +#define TENT_PLATFORM_TPU_H_ + +#include "tent/common/config.h" +#include "tent/platform/cpu.h" +#include "tent/runtime/platform.h" + +namespace mooncake { +namespace tent { + +// TpuPlatform models a TPU host: the host DRAM and RDMA topology are identical +// to CpuPlatform, so those paths (host allocation, NUMA probing, host<->host +// copy) are inherited unchanged. Only the TPU-device-aware operations are +// overridden, and every one of them is delegated to TpuPjrtShim so that this +// class carries no direct PJRT dependency. +// +// TPU HBM cannot be reached by the NIC, so there is no direct device transport: +// the HBM<->host hop is performed here via copy(), and the host<->host hop is +// carried by RDMA/TCP. ProxyManager chains the two (see findStagingPolicy). +class TpuPlatform : public CpuPlatform { + public: + explicit TpuPlatform(std::shared_ptr config) + : CpuPlatform(config) {} + + ~TpuPlatform() override {} + + // Host + RDMA discovery from CpuPlatform, plus one MemEntry per TPU device + // ("tpu:N") so findNearMem() can resolve a device to its nearest host node. + Status probe(std::vector &nic_list, + std::vector &mem_list) override; + + // Device allocation is owned by the serving framework (JAX / torch-XLA); + // host allocation is inherited. "tpu" locations therefore return + // NotImplemented here. + Status allocate(void **pptr, size_t size, MemoryOptions &options) override; + + // HBM<->host copy via the adapter when either side is TPU memory; otherwise + // the inherited host memcpy. + Status copy(void *dst, void *src, size_t length) override; + + MemoryType getMemoryType(void *addr) override; + + const std::vector getLocation( + void *start, size_t len, bool skip_prefault = false) override; + + const std::string type() const override { return "tpu"; } +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_PLATFORM_TPU_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h new file mode 100644 index 0000000000..ceb7669785 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h @@ -0,0 +1,64 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// C ABI contract between TENT and the TPU/PJRT adapter shared library. +// +// The adapter is built separately against the PJRT runtime and exports these +// symbols with C linkage. TENT resolves them at runtime via dlopen()/dlsym() +// (see TpuPjrtShim); it never links the adapter or the PJRT runtime directly. +// +// Pointer classification note: the "device pointer" passed across this ABI is +// the stable token the serving-engine integration registers with TENT for a TPU +// buffer (see registerLocalMemory with a "tpu:N" location). The adapter owns +// the mapping from that token to the underlying PJRT buffer; mc_tpu_pjrt_* copy +// and classification calls resolve the token through the adapter's own +// registry. + +#ifndef TENT_PLATFORM_TPU_PJRT_ABI_H_ +#define TENT_PLATFORM_TPU_PJRT_ABI_H_ + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Initialize the adapter (creates/attaches the PJRT client). Returns 0 on +// success, non-zero on failure. Idempotent; safe to call more than once. +int mc_tpu_pjrt_init(void); + +// Returns 1 if `addr` is a TPU device buffer known to the adapter, else 0. +int mc_tpu_pjrt_is_device_ptr(const void *addr); + +// Returns the device ordinal backing `addr`, or -1 if `addr` is not a known TPU +// device buffer. +int mc_tpu_pjrt_device_index(const void *addr); + +// Synchronous device->host copy. Returns 0 on success, non-zero on failure. +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len); + +// Synchronous host->device copy. Returns 0 on success, non-zero on failure. +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len); + +// Number of visible TPU devices. +int mc_tpu_pjrt_device_count(void); + +// NUMA node closest to device `index`, or -1 if unknown. +int mc_tpu_pjrt_device_numa(int index); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TENT_PLATFORM_TPU_PJRT_ABI_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h new file mode 100644 index 0000000000..0d26595c7b --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h @@ -0,0 +1,93 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_PLATFORM_TPU_PJRT_SHIM_H_ +#define TENT_PLATFORM_TPU_PJRT_SHIM_H_ + +#include +#include + +#include "tent/common/status.h" + +namespace mooncake { +namespace tent { + +// TpuPjrtShim isolates every TPU/PJRT dependency behind a narrow interface. +// +// The TPU device I/O primitives (HBM<->host DMA, device-pointer classification, +// device topology) are provided by a separate adapter shared library built +// against the PJRT runtime. TENT itself carries no build-time dependency on +// that runtime: the adapter is resolved at runtime via dlopen(), matching the +// way the other accelerator backends keep vendor SDKs out of the core build. +// +// The adapter library must export the C ABI declared in tpu_pjrt_abi.h. Its +// path defaults to "libmooncake_tpu_pjrt.so" and can be overridden with the +// MC_TPU_PJRT_LIB environment variable. When the adapter cannot be loaded, +// available() returns false and every operation returns a non-OK Status; this +// keeps a USE_TPU build functional (and unit-testable with a mock adapter) +// without the real runtime present. +class TpuPjrtShim { + public: + // Process-wide singleton. The adapter is loaded (and initialized) lazily on + // first use; loading is attempted at most once. + static TpuPjrtShim &instance(); + + // True when the adapter library was loaded and initialized successfully. + bool available() const { return available_; } + + // Returns true if `addr` refers to memory owned by the TPU runtime (HBM). + // Returns false when the adapter is unavailable or the pointer is host + // memory, so a caller can safely treat "not TPU" as host memory. + bool isDevicePtr(const void *addr) const; + + // Device ordinal backing `addr`, or -1 if `addr` is not TPU device memory. + int deviceIndex(const void *addr) const; + + // Synchronous HBM -> host DMA copy of `length` bytes. + Status copyD2H(void *host_dst, const void *device_src, size_t length) const; + + // Synchronous host -> HBM DMA copy of `length` bytes. + Status copyH2D(void *device_dst, const void *host_src, size_t length) const; + + // Number of visible TPU devices (0 when the adapter is unavailable). + int deviceCount() const; + + // NUMA node closest to TPU device `index`, or -1 if unknown. + int deviceNumaNode(int index) const; + + private: + TpuPjrtShim(); + ~TpuPjrtShim(); + TpuPjrtShim(const TpuPjrtShim &) = delete; + TpuPjrtShim &operator=(const TpuPjrtShim &) = delete; + + void load(); + + void *handle_ = nullptr; + bool available_ = false; + + // Resolved adapter entrypoints (see tpu_pjrt_abi.h for the contract). + int (*fn_init_)() = nullptr; + int (*fn_is_device_ptr_)(const void *) = nullptr; + int (*fn_device_index_)(const void *) = nullptr; + int (*fn_copy_d2h_)(void *, const void *, size_t) = nullptr; + int (*fn_copy_h2d_)(void *, const void *, size_t) = nullptr; + int (*fn_device_count_)() = nullptr; + int (*fn_device_numa_)(int) = nullptr; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_PLATFORM_TPU_PJRT_SHIM_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h index 95a4c7aefa..e53265adb0 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h @@ -20,7 +20,10 @@ namespace mooncake { namespace tent { -enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA, MTYPE_ROCM }; +// MTYPE_TPU is appended last so the numeric values of the existing entries are +// preserved. TPU HBM is not NIC-addressable, so transfers touching it are +// staged through host DRAM by ProxyManager (see findStagingPolicy). +enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA, MTYPE_ROCM, MTYPE_TPU }; class Platform { public: diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h new file mode 100644 index 0000000000..e4d0890d4c --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h @@ -0,0 +1,93 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_TPU_TRANSPORT_H_ +#define TENT_TPU_TRANSPORT_H_ + +#include +#include + +#include "tent/runtime/control_plane.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { + +struct TpuTask { + Request request; + volatile TransferStatusEnum status_word; + volatile size_t transferred_bytes; +}; + +struct TpuSubBatch : public Transport::SubBatch { + std::vector task_list; + size_t max_size; + virtual size_t size() const { return task_list.size(); } +}; + +// TpuTransport is the local staging executor for TPU: it performs the +// HBM<->host-DRAM hop of a staged transfer by delegating to +// Platform::copy() (which routes to the PJRT adapter for TPU memory). It never +// touches the network — the host<->host hop is carried by RDMA/TCP, and +// ProxyManager chains the two stages (see findStagingPolicy). +// +// Because TPU HBM is not NIC-addressable, this transport advertises only the +// device<->host capabilities (gpu_to_dram / dram_to_gpu) and leaves gpu_to_gpu +// false so the engine always stages cross-node traffic through host DRAM. It is +// therefore only ever selected for LOCAL_SEGMENT_ID copies. +class TpuTransport : public Transport { + public: + TpuTransport(); + + ~TpuTransport(); + + virtual Status install(std::string &local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr); + + virtual Status uninstall(); + + virtual Status allocateSubBatch(SubBatchRef &batch, size_t max_size); + + virtual Status freeSubBatch(SubBatchRef &batch); + + virtual Status submitTransferTasks( + SubBatchRef batch, const std::vector &request_list); + + virtual Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus &status); + + virtual Status addMemoryBuffer(BufferDesc &desc, + const MemoryOptions &options); + + virtual Status removeMemoryBuffer(BufferDesc &desc); + + virtual const char *getName() const { return "tpu"; } + + private: + void startTransfer(TpuTask *task, TpuSubBatch *batch); + + private: + bool installed_; + std::string local_segment_name_; + std::shared_ptr local_topology_; + std::shared_ptr metadata_; + std::shared_ptr conf_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_TPU_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/src/CMakeLists.txt b/mooncake-transfer-engine/tent/src/CMakeLists.txt index 73766ce120..0b7b5a4afc 100644 --- a/mooncake-transfer-engine/tent/src/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/CMakeLists.txt @@ -131,6 +131,7 @@ foreach( platform_rocm platform_ascend platform_sunrise + platform_tpu tent_xport_gds tent_xport_uring tent_xport_bufio @@ -141,6 +142,7 @@ foreach( tent_xport_tcp tent_xport_ascend_direct tent_xport_sunrise_link + tent_xport_tpu tent_metrics) if(TARGET ${tgt}) target_link_libraries(tent_link_group INTERFACE ${tgt}) diff --git a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt index 2f48b68696..11b7bbd9ac 100644 --- a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt @@ -2,6 +2,7 @@ add_subdirectory(cuda) add_subdirectory(rocm) add_subdirectory(ascend) add_subdirectory(sunrise) +add_subdirectory(tpu) file(GLOB PLATFORM_SOURCES "*.cpp") add_library(tent_platform_all STATIC ${PLATFORM_SOURCES}) @@ -18,3 +19,6 @@ endif() if(TARGET platform_sunrise) target_link_libraries(tent_platform_all PUBLIC platform_sunrise) endif() +if(TARGET platform_tpu) + target_link_libraries(tent_platform_all PUBLIC platform_tpu) +endif() diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt new file mode 100644 index 0000000000..c59d55d480 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt @@ -0,0 +1,7 @@ +if(USE_TPU) + file(GLOB TENT_PLATFORM_TPU_SOURCES "*.cpp") + add_library(platform_tpu STATIC ${TENT_PLATFORM_TPU_SOURCES}) + # No PJRT / TPU SDK link dependency: the adapter is resolved at runtime via + # dlopen() in TpuPjrtShim, so we only need libdl. + target_link_libraries(platform_tpu PUBLIC tent_common ${CMAKE_DL_LIBS}) +endif() diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/README.md b/mooncake-transfer-engine/tent/src/platform/tpu/README.md new file mode 100644 index 0000000000..71f3cd4483 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/README.md @@ -0,0 +1,73 @@ +# TPU (PJRT) platform for TENT + +This directory implements Google TPU support for the TENT transfer engine. +It is compiled only when TENT is built with `-DUSE_TPU=ON` (OFF by default). + +## Why staging + +TPU HBM is not addressable by the NIC, so a transfer that touches TPU memory +cannot be issued directly by RDMA/TCP. Instead every such transfer is **staged +through host DRAM** and the two hops are chained by the existing +`ProxyManager` pipeline: + +``` + TPU HBM <-> host DRAM (PJRT device copy, this platform) + host DRAM <-> remote host DRAM (RDMA / TCP, existing transports) + remote host DRAM <-> remote TPU HBM (PJRT device copy on the peer) +``` + +No new networked transport is required — TENT reuses `ProxyManager`'s chunked +double-buffering, staging-buffer lifecycle, and async status tracking. + +> New to the codebase and wondering why this needs a `TpuTransport` at all +> instead of a one-line branch in `ProxyManager`? See +> [`STAGING_ARCHITECTURE.md`](STAGING_ARCHITECTURE.md). + +## Components + +| File | Role | +|------|------| +| `tpu_platform.cpp` (`TpuPlatform`) | Derives `CpuPlatform`. Host DRAM + NIC topology are inherited (RDMA NICs if present, otherwise just host NUMA nodes for TCP); only the TPU-device-aware paths are overridden (`copy`, `getMemoryType`, `getLocation`, and one `MemEntry` per device in `probe`). | +| `tpu_pjrt_shim.cpp` (`TpuPjrtShim`) | Isolates all PJRT dependency. Resolves the device-copy adapter at runtime via `dlopen`. | +| `../../transport/tpu/tpu_transport.cpp` (`TpuTransport`) | The local HBM↔host staging executor. Advertises only `gpu_to_dram` / `dram_to_gpu`; runs the copy via `Platform::copy` for `LOCAL_SEGMENT_ID` requests. | + +The staging policy that ties these together lives in +`TransferEngineImpl::findStagingPolicy` (case "TPU"). + +## The device-copy adapter (runtime dependency) + +`TpuPjrtShim` does **not** link the PJRT runtime. Instead it loads an adapter +shared library at runtime that exports the C ABI declared in +[`tpu_pjrt_abi.h`](../../../include/tent/platform/tpu_pjrt_abi.h): + +```c +int mc_tpu_pjrt_init(void); +int mc_tpu_pjrt_is_device_ptr(const void *addr); +int mc_tpu_pjrt_device_index(const void *addr); +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len); +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len); +int mc_tpu_pjrt_device_count(void); +int mc_tpu_pjrt_device_numa(int index); +``` + +- **Discovery:** the library path defaults to `libmooncake_tpu_pjrt.so` and can + be overridden with the `MC_TPU_PJRT_LIB` environment variable. +- **Graceful absence:** if the adapter cannot be loaded or does not satisfy the + ABI, `TpuPjrtShim::available()` returns false and all TPU operations return a + non-OK `Status`. A `USE_TPU` build therefore links and runs without the + runtime present (useful for CI and unit tests). +- **Pointer tokens:** the `const void *` "device pointer" is the stable token the + serving-engine integration registers with TENT for a TPU buffer (via a + `tpu:N` location). The adapter owns the mapping from that token to the + underlying PJRT buffer. + +A mock adapter and a unit test that exercise this ABI on any Linux host live in +[`../../../tests/tpu/`](../../../tests/tpu/). + +## Not yet included (follow-ups) + +- Serving-framework (JAX / PyTorch-XLA) integration that registers TPU buffers. +- DMA-mapped (pinned) staging buffers for true async device DMA, and the + interaction between device DMA-mapping and RDMA memory registration on the + same host buffer. +- Benchmarks on real TPU hardware. diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md b/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md new file mode 100644 index 0000000000..119909b9de --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md @@ -0,0 +1,171 @@ +# Why TPU support needs a `TpuTransport` (and not just an `if (pjrt)` branch) + +This note explains, for someone new to the TENT codebase, **why adding TPU +staging touches a transport class, a platform class, and the routing/capability +system** rather than being a one-line change inside `ProxyManager`. Everything +here is about *existing* TENT mechanics; TPU is just the motivating example. + +All paths below are under `mooncake-transfer-engine/tent/`. + +--- + +## 1. The three layers you need to know + +TENT separates "how do I move bytes" into three cooperating layers: + +| Layer | What it is | Example | +|-------|-----------|---------| +| **Platform** | The *local* memory primitives for one accelerator family: allocate, free, `copy`, classify a pointer's memory type, discover topology. One process picks exactly one Platform at build time. | `CpuPlatform`, `CudaPlatform`, `TpuPlatform` | +| **Transport** | A *data channel* that moves a transfer request's bytes. Each advertises **capabilities** (can it do dram→dram? gpu→dram? gpu→gpu?). | `RdmaTransport`, `ShmTransport`, `NvlinkTransport`, `TpuTransport` | +| **Staging (ProxyManager)** | When no single transport can do a hop directly, it splits the transfer into stages through host DRAM and chains transports. | `ProxyManager` | + +Key headers: `include/tent/runtime/platform.h`, `include/tent/runtime/transport.h` +(the `Capabilities` struct is at the top of the latter). + +--- + +## 2. The naive expectation + +> "TPU HBM can't be reached by the NIC, so stage it through host DRAM. TENT +> already stages CUDA that way, so just find where `ProxyManager` does the +> device copy and add an `if (tpu) pjrt_copy() else cudaMemcpy()`." + +That mental model is *almost* right about the data flow but wrong about the +mechanism. There is no `cudaMemcpy` inside `ProxyManager` to branch on. + +--- + +## 3. What actually happens when you submit a transfer + +Follow one transfer from the top: + +1. **Submit** → `TransferEngineImpl::submitTransfer` → `prepareSubmit` + (`src/runtime/transfer_engine_impl.cpp`). + +2. **Should this be staged?** For each request, `prepareSubmit` calls + [`findStagingPolicy`](../../runtime/transfer_engine_impl.cpp) (defined at + `transfer_engine_impl.cpp:1278`). It returns a 3-element plan + `[server, local_stage_location, remote_stage_location]`; empty entries mean + "no staging on that side". Then: + ```cpp + owner.staging = !owner.staging_params.empty() && staging_proxy_; // ~line 1401 + ``` + +3. **If staged**, the task is handed to the proxy instead of a transport: + ```cpp + staging_proxy_->submit(&task, batch, owner.staging_params); // ~line 1477 + ``` + +4. **`ProxyManager` chops the transfer into chunks** and, per chunk, issues + *sub-transfers* for each stage — it does **not** copy bytes itself. See + `ProxyManager::transferEventLoop` (`src/runtime/proxy_manager.cpp:251`) and + `submitLocalStage` (`:71`). The local stage is submitted as an ordinary + transfer whose target is `LOCAL_SEGMENT_ID`: + ```cpp + // local_stage.source = device pointer, target_offset = host staging buffer + impl_->submitStagingTransfer(batch, {local_stage}); // -> submitTransfer(...) + ``` + +5. **That sub-transfer is routed like any other**, back through + `getTransportType`, and the *selected transport's* `submitTransferTasks` + runs. For the simple copy transports, that method is where the actual byte + movement lives — and it calls **`Platform::copy`**, not any device API + directly. Example, `ShmTransport::startTransfer` + (`src/transport/shm/shm_transport.cpp:119`): + ```cpp + status = Platform::getLoader().copy(dst, src, length); // :122 / :126 + ``` + +So the real device-copy seam is **`Platform::copy`** +(`CudaPlatform::copy` does `cudaMemcpyAsync`; `TpuPlatform::copy` will call the +PJRT adapter). That part *is* a clean override — good. + +**But there is a gate before you ever reach step 5.** + +--- + +## 4. The crux: routing is gated by transport *capabilities* + +In step 5 the sub-transfer must be *routed to some transport*. Routing asks each +candidate transport: "can you do a copy from `local_memory_type` to +`remote_memory_type`?" via a capability check: + +- Selector mode (**the default**): `TransportSelector::isTransportAvailable` + (`src/runtime/transport_selector.cpp:311`). +- Legacy mode: `checkAvailability` (`src/runtime/transfer_engine_impl.cpp:889`). + +Both map a `(local, remote)` memory-type pair to one capability bit: + +```cpp +// GPU/device local side, CPU peer -> needs caps.gpu_to_dram +// CPU local side, GPU/device peer -> needs caps.dram_to_gpu +// GPU to GPU -> needs caps.gpu_to_gpu +// CPU to CPU -> needs caps.dram_to_dram +``` + +For the TPU **local stage**, the sub-transfer is `TPU device -> host DRAM`, so +routing looks for a transport advertising `gpu_to_dram`. + +Now look at who advertises what: + +| Transport | Advertises | Set where | +|-----------|-----------|-----------| +| `RdmaTransport` | `dram_to_dram`; adds `gpu_*` **only if `nvidia_peermem` is loaded** (CUDA GPUDirect) | `rdma_transport.cpp` `install` | +| `ShmTransport` | `dram_to_dram` **only** | `shm_transport.cpp` `install` | +| `NvlinkTransport` | `dram_to_gpu` / `gpu_to_dram` / `gpu_to_gpu` (CUDA) | `nvlink_transport.cpp` `install` | + +On a TPU host there is **no transport that advertises `gpu_to_dram`**: +- `ShmTransport` = `dram_to_dram` only. +- `RdmaTransport` won't set `gpu_to_dram` (no `nvidia_peermem`), and it *couldn't* + DMA out of HBM anyway. + +**Therefore, without a new transport, the TPU local-stage sub-transfer has no +route, and it fails — even though `TpuPlatform::copy` exists and could do the +copy.** `findStagingPolicy` would also have nothing to hang the policy on. + +This is the answer to "why not just one branch": the copy *mechanism* is a clean +`Platform::copy` override, but the *routing/capability system* has no way to +select it unless some `Transport` object sits in `transport_list_[...]` and +answers "yes, I can do `gpu_to_dram`." + +For CUDA, that answering transport is `NvlinkTransport` — a full P2P transport +that TENT reuses. TPU has no equivalent to piggyback on, so we add a **thin one +whose only job is to advertise the device↔host capability and run +`Platform::copy` for `LOCAL_SEGMENT_ID` requests**: `TpuTransport` +(`src/transport/tpu/tpu_transport.cpp`). It never touches the network; the +host↔host hop is still RDMA/TCP. + +--- + +## 5. What PR #1 therefore has to touch + +| Concern | Change | File | +|---------|--------|------| +| A memory type for TPU HBM | `MTYPE_TPU` + `"tpu"` parsing | `platform.h`, `getTypeEnum` in `transfer_engine_impl.cpp` | +| Local device copy primitive | `TpuPlatform::copy` → PJRT adapter (the clean seam) | `src/platform/tpu/tpu_platform.cpp` | +| A transport that *advertises* `gpu_to_dram`/`dram_to_gpu` so routing can pick the copy | thin `TpuTransport` | `src/transport/tpu/tpu_transport.cpp` | +| Teach the capability checks that TPU is a device type | add `MTYPE_TPU` to `isGpuType` **and** the selector's `is_gpu` lambda | `transfer_engine_impl.cpp:880`, `transport_selector.cpp:337` | +| Decide to stage TPU transfers | `findStagingPolicy` TPU case | `transfer_engine_impl.cpp:1278` | + +Note the **two** capability helpers (line 880 and 337): there are two routing +modes (selector = default, legacy), each with its own device-type predicate. +Both must learn about `MTYPE_TPU`, or TPU works in one mode and silently fails in +the other. This is exactly the kind of thing that isn't visible until you trace +both paths. + +--- + +## 6. One-paragraph summary + +TENT's `ProxyManager` already implements the staged-DRAM pipeline, and +`Platform::copy` is a clean place to plug in a PJRT device copy. But a staged +sub-transfer is still *routed*, and routing only selects a transport that +*advertises* the matching capability. No existing transport advertises +`device↔host` on a TPU host, so the copy would never be reached. `TpuTransport` +exists solely to answer that capability query and dispatch to `Platform::copy`; +it is the unavoidable "interface tax" of participating in TENT's transport +system, and it accounts for most of the non-test, non-doc line count in this +change. + +See also: [`README.md`](README.md) in this directory for the component list and +the PJRT adapter ABI. diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp new file mode 100644 index 0000000000..164dc13bae --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp @@ -0,0 +1,144 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/platform/tpu_pjrt_shim.h" + +#include +#include + +#include + +namespace mooncake { +namespace tent { + +namespace { +constexpr const char kDefaultAdapterLib[] = "libmooncake_tpu_pjrt.so"; + +const char *adapterLibraryPath() { + const char *env = std::getenv("MC_TPU_PJRT_LIB"); + if (env && env[0] != '\0') return env; + return kDefaultAdapterLib; +} +} // namespace + +TpuPjrtShim &TpuPjrtShim::instance() { + static TpuPjrtShim g_instance; + return g_instance; +} + +TpuPjrtShim::TpuPjrtShim() { load(); } + +TpuPjrtShim::~TpuPjrtShim() { + if (handle_) dlclose(handle_); +} + +void TpuPjrtShim::load() { + const char *lib = adapterLibraryPath(); + handle_ = dlopen(lib, RTLD_NOW | RTLD_LOCAL); + if (!handle_) { + LOG(WARNING) << "TpuPjrtShim: unable to load TPU PJRT adapter '" << lib + << "': " << dlerror() + << ". TPU transfers will be unavailable. Set " + "MC_TPU_PJRT_LIB to override the adapter path."; + available_ = false; + return; + } + + // Resolve every entrypoint; treat any missing symbol as a fatal load error + // so we never partially bind an incompatible adapter. + auto resolve = [&](const char *name) -> void * { + void *sym = dlsym(handle_, name); + if (!sym) + LOG(WARNING) << "TpuPjrtShim: adapter '" << lib + << "' is missing symbol '" << name << "'"; + return sym; + }; + + fn_init_ = reinterpret_cast(resolve("mc_tpu_pjrt_init")); + fn_is_device_ptr_ = reinterpret_cast( + resolve("mc_tpu_pjrt_is_device_ptr")); + fn_device_index_ = reinterpret_cast( + resolve("mc_tpu_pjrt_device_index")); + fn_copy_d2h_ = reinterpret_cast( + resolve("mc_tpu_pjrt_copy_d2h")); + fn_copy_h2d_ = reinterpret_cast( + resolve("mc_tpu_pjrt_copy_h2d")); + fn_device_count_ = + reinterpret_cast(resolve("mc_tpu_pjrt_device_count")); + fn_device_numa_ = + reinterpret_cast(resolve("mc_tpu_pjrt_device_numa")); + + if (!fn_init_ || !fn_is_device_ptr_ || !fn_device_index_ || !fn_copy_d2h_ || + !fn_copy_h2d_ || !fn_device_count_ || !fn_device_numa_) { + LOG(ERROR) << "TpuPjrtShim: adapter '" << lib + << "' does not satisfy the required ABI; disabling TPU."; + dlclose(handle_); + handle_ = nullptr; + available_ = false; + return; + } + + if (fn_init_() != 0) { + LOG(ERROR) << "TpuPjrtShim: mc_tpu_pjrt_init() failed; disabling TPU."; + dlclose(handle_); + handle_ = nullptr; + available_ = false; + return; + } + + available_ = true; + LOG(INFO) << "TpuPjrtShim: TPU PJRT adapter '" << lib << "' loaded (" + << fn_device_count_() << " device(s))."; +} + +bool TpuPjrtShim::isDevicePtr(const void *addr) const { + if (!available_ || !addr) return false; + return fn_is_device_ptr_(addr) != 0; +} + +int TpuPjrtShim::deviceIndex(const void *addr) const { + if (!available_ || !addr) return -1; + return fn_device_index_(addr); +} + +Status TpuPjrtShim::copyD2H(void *host_dst, const void *device_src, + size_t length) const { + if (!available_) + return Status::NotImplemented("TPU PJRT adapter unavailable" LOC_MARK); + if (fn_copy_d2h_(host_dst, device_src, length) != 0) + return Status::InternalError("TPU device->host copy failed" LOC_MARK); + return Status::OK(); +} + +Status TpuPjrtShim::copyH2D(void *device_dst, const void *host_src, + size_t length) const { + if (!available_) + return Status::NotImplemented("TPU PJRT adapter unavailable" LOC_MARK); + if (fn_copy_h2d_(device_dst, host_src, length) != 0) + return Status::InternalError("TPU host->device copy failed" LOC_MARK); + return Status::OK(); +} + +int TpuPjrtShim::deviceCount() const { + if (!available_) return 0; + return fn_device_count_(); +} + +int TpuPjrtShim::deviceNumaNode(int index) const { + if (!available_) return -1; + return fn_device_numa_(index); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp new file mode 100644 index 0000000000..625ea8edd1 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp @@ -0,0 +1,103 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/platform/tpu.h" + +#include + +#include + +#include "tent/common/status.h" +#include "tent/platform/tpu_pjrt_shim.h" + +namespace mooncake { +namespace tent { + +Status TpuPlatform::probe(std::vector& nic_list, + std::vector& mem_list) { + // Host DRAM + RDMA topology is identical to a CPU host. + CHECK_STATUS(CpuPlatform::probe(nic_list, mem_list)); + + // Register one memory node per visible TPU device so findNearMem("tpu:N") + // resolves to the nearest host NUMA node for staging. Device I/O never + // touches the NIC directly, so these entries only carry NIC affinity used + // to place the host staging buffers. + auto& shim = TpuPjrtShim::instance(); + int device_count = shim.deviceCount(); + for (int i = 0; i < device_count; ++i) { + Topology::MemEntry entry; + entry.name = "tpu:" + std::to_string(i); + entry.numa_node = shim.deviceNumaNode(i); + entry.type = Topology::MEM_UNKNOWN; + int nic_id = 0; + for (const auto& nic : nic_list) { + if (entry.numa_node >= 0 && nic.numa_node == entry.numa_node) + entry.device_list[0].push_back(nic_id); + else + entry.device_list[2].push_back(nic_id); + nic_id++; + } + mem_list.push_back(std::move(entry)); + } + return Status::OK(); +} + +Status TpuPlatform::allocate(void** pptr, size_t size, MemoryOptions& options) { + LocationParser location(options.location); + if (location.type() == "tpu") { + // TPU HBM buffers are owned by the serving framework and registered + // with TENT; TENT never allocates them itself. + return Status::NotImplemented( + "TpuPlatform does not allocate TPU device memory" LOC_MARK); + } + // Host DRAM staging buffers use the inherited NUMA-aware allocator. + return CpuPlatform::allocate(pptr, size, options); +} + +Status TpuPlatform::copy(void* dst, void* src, size_t length) { + auto& shim = TpuPjrtShim::instance(); + bool src_is_device = shim.isDevicePtr(src); + bool dst_is_device = shim.isDevicePtr(dst); + if (src_is_device && dst_is_device) { + return Status::NotImplemented( + "TpuPlatform: device-to-device copy is not supported; transfers " + "are staged through host DRAM" LOC_MARK); + } + if (src_is_device) return shim.copyD2H(dst, src, length); + if (dst_is_device) return shim.copyH2D(dst, src, length); + // Neither side is TPU memory: plain host copy. + return CpuPlatform::copy(dst, src, length); +} + +MemoryType TpuPlatform::getMemoryType(void* addr) { + if (TpuPjrtShim::instance().isDevicePtr(addr)) return MTYPE_TPU; + return CpuPlatform::getMemoryType(addr); +} + +const std::vector TpuPlatform::getLocation(void* start, + size_t len, + bool skip_prefault) { + auto& shim = TpuPjrtShim::instance(); + if (shim.isDevicePtr(start)) { + int index = shim.deviceIndex(start); + std::string location = + index >= 0 ? "tpu:" + std::to_string(index) : kWildcardLocation; + return {RangeLocation{reinterpret_cast(start), len, + std::move(location)}}; + } + return CpuPlatform::getLocation(start, len, skip_prefault); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/runtime/platform.cpp b/mooncake-transfer-engine/tent/src/runtime/platform.cpp index 72a9d096d9..611737f532 100644 --- a/mooncake-transfer-engine/tent/src/runtime/platform.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/platform.cpp @@ -22,6 +22,8 @@ #include "tent/platform/sunrise.h" #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) #include "tent/platform/ascend.h" +#elif defined(USE_TPU) +#include "tent/platform/tpu.h" #else #include "tent/platform/cpu.h" #endif @@ -41,6 +43,8 @@ Platform& Platform::getLoader(std::shared_ptr conf) { g_instance = std::make_shared(conf); #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) g_instance = std::make_shared(conf); +#elif defined(USE_TPU) + g_instance = std::make_shared(conf); #else g_instance = std::make_shared(conf); #endif diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 2e3f345540..b1f2116eee 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -645,6 +645,7 @@ std::vector TransferEngineImpl::getSupportedTransports( if (transport_list_[SHM]) result.push_back(SHM); if (transport_list_[TCP]) result.push_back(TCP); if (transport_list_[GDS]) result.push_back(GDS); + if (transport_list_[TPU]) result.push_back(TPU); return result; } @@ -877,7 +878,12 @@ class TransferEngineImpl::BatchRef { }; static bool isGpuType(MemoryType t) { - return t == MTYPE_CUDA || t == MTYPE_ROCM; + // TPU HBM behaves like a GPU that lacks NIC access: it is a device-side + // memory that can only reach the network by staging through host DRAM. + // Treating it as a "gpu type" makes the capability checks route its + // device<->host hop to TpuTransport (gpu_to_dram / dram_to_gpu) while + // leaving gpu_to_gpu unsatisfiable, which forces host-DRAM staging. + return t == MTYPE_CUDA || t == MTYPE_ROCM || t == MTYPE_TPU; } static bool checkAvailability(const std::shared_ptr& xport, @@ -905,6 +911,7 @@ static MemoryType getTypeEnum(const std::string& type) { if (type == "cuda") return MTYPE_CUDA; if (type == "npu") return MTYPE_CUDA; if (type == "rocm") return MTYPE_ROCM; + if (type == "tpu") return MTYPE_TPU; return MTYPE_UNKNOWN; } @@ -964,7 +971,11 @@ SelectionResult TransferEngineImpl::getTransportType(const Request& request, auto remote_mtype = getTypeEnum(LocationParser(entry->location).type()); for (auto type : entry->transports) { - if ((type == NVLINK || type == SHM) && !same_machine) + // NVLINK/SHM are same-machine only; TPU is a + // local-stage-only executor and must never carry a remote + // hop. + if ((type == NVLINK || type == SHM || type == TPU) && + !same_machine) continue; if (checkAvailability(transport_list_[type], local_mtype, remote_mtype)) { @@ -1042,6 +1053,8 @@ static const char* transportTypeName(TransportType type) { return "AscendDirect"; case SUNRISE_LINK: return "SUNRISE_LINK"; + case TPU: + return "TPU"; } return "UNKNOWN"; } @@ -1328,6 +1341,34 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, remote, Topology::MEM_CUDA)); } } + // case 3: TPU. HBM is not NIC-addressable, so any hop touching TPU memory + // is staged through host DRAM: TpuTransport performs the local HBM<->host + // copy (via the PJRT adapter) and the host<->host hop is carried by + // whatever host-DRAM network transport is present. TPU deployments (e.g. + // cloud TPU VMs) are typically TCP/multi-NIC rather than RDMA, so we gate + // on either; the cross stage itself is routed by capability (dram_to_dram), + // so TCP is selected when RDMA is absent. We also require TpuTransport (the + // local HBM<->host executor), mirroring how the CUDA cases gate on NVLINK. + // An empty stage location means "no staging needed on that side". + if (transport_list_[TPU] && + (transport_list_[RDMA] || transport_list_[TCP])) { + if (local_mtype == MTYPE_TPU && remote_mtype == MTYPE_TPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(topology_->findNearMem(local)); + policy.push_back(desc->getMemory().topology.findNearMem(remote)); + } else if (local_mtype == MTYPE_TPU && remote_mtype == MTYPE_CPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(topology_->findNearMem(local)); + policy.push_back(""); // remote already host DRAM + } else if (local_mtype == MTYPE_CPU && remote_mtype == MTYPE_TPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(""); // local already host DRAM + policy.push_back(desc->getMemory().topology.findNearMem(remote)); + } + } } SelectionResult TransferEngineImpl::resolveTransport(const Request& req, diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp index 7ecd51276e..6f3a2584b8 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp @@ -41,6 +41,10 @@ #include "tent/transport/sunrise_link/sunrise_link_transport.h" #endif +#ifdef USE_TPU +#include "tent/transport/tpu/tpu_transport.h" +#endif + namespace mooncake { namespace tent { @@ -94,6 +98,11 @@ Status TransferEngineImpl::loadTransports() { } #endif +#ifdef USE_TPU + if (conf_->get("transports/tpu/enable", true)) + transport_list_[TPU] = std::make_shared(); +#endif + return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp index e26f1e78a0..fc53df32cf 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp @@ -29,20 +29,32 @@ namespace tent { // Transport type name mapping static const std::unordered_map kTransportNameMap = { - {"unspec", UNSPEC}, {"rdma", RDMA}, - {"mnnvl", MNNVL}, {"shm", SHM}, - {"nvlink", NVLINK}, {"gds", GDS}, - {"io_uring", IOURING}, {"tcp", TCP}, - {"ascend", AscendDirect}, {"sunrise_link", SUNRISE_LINK}, + {"unspec", UNSPEC}, + {"rdma", RDMA}, + {"mnnvl", MNNVL}, + {"shm", SHM}, + {"nvlink", NVLINK}, + {"gds", GDS}, + {"io_uring", IOURING}, + {"tcp", TCP}, + {"ascend", AscendDirect}, + {"sunrise_link", SUNRISE_LINK}, + {"tpu", TPU}, }; static const std::unordered_map kTransportTypeNames = { - {UNSPEC, "unspec"}, {RDMA, "rdma"}, - {MNNVL, "mnnvl"}, {SHM, "shm"}, - {NVLINK, "nvlink"}, {GDS, "gds"}, - {IOURING, "io_uring"}, {TCP, "tcp"}, - {AscendDirect, "ascend"}, {SUNRISE_LINK, "sunrise_link"}, + {UNSPEC, "unspec"}, + {RDMA, "rdma"}, + {MNNVL, "mnnvl"}, + {SHM, "shm"}, + {NVLINK, "nvlink"}, + {GDS, "gds"}, + {IOURING, "io_uring"}, + {TCP, "tcp"}, + {AscendDirect, "ascend"}, + {SUNRISE_LINK, "sunrise_link"}, + {TPU, "tpu"}, }; // Memory type name mapping for pattern matching @@ -236,6 +248,9 @@ bool TransportSelector::matchesMemoryPattern(const std::string& pattern, case MTYPE_ROCM: type_str = "rocm"; break; + case MTYPE_TPU: + type_str = "tpu"; + break; default: type_str = "unknown"; break; @@ -319,15 +334,21 @@ bool TransportSelector::isTransportAvailable( } // Special constraints - if ((type == NVLINK || type == SHM) && !context.same_machine) { - return false; // NVLINK and SHM only work on same machine + if ((type == NVLINK || type == SHM || type == TPU) && + !context.same_machine) { + // NVLINK/SHM only work on same machine; TPU is a local-stage-only + // executor (HBM<->host), so it must never be picked for a remote hop. + return false; } const auto& caps = transport->capabilities(); - // Helper to check if memory type is GPU/NPU + // Helper to check if memory type is a device (GPU/NPU/TPU). TPU is included + // so its device<->host staging hop routes to TpuTransport (gpu_to_dram / + // dram_to_gpu); it never satisfies gpu_to_gpu, so cross-node TPU traffic is + // always staged through host DRAM. auto is_gpu = [](MemoryType t) { - return t == MTYPE_CUDA || t == MTYPE_ROCM; + return t == MTYPE_CUDA || t == MTYPE_ROCM || t == MTYPE_TPU; }; // For file segments, check file-specific capabilities (original logic) diff --git a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt index e73f7cb063..05140d6987 100644 --- a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt @@ -8,6 +8,7 @@ add_subdirectory(io_uring) add_subdirectory(bufio) add_subdirectory(ascend) add_subdirectory(sunrise_link) +add_subdirectory(tpu) add_library(tent_transport_all INTERFACE) foreach( @@ -21,7 +22,8 @@ foreach( tent_xport_shm tent_xport_tcp tent_xport_ascend_direct - tent_xport_sunrise_link) + tent_xport_sunrise_link + tent_xport_tpu) if(TARGET ${tgt}) target_link_libraries(tent_transport_all INTERFACE ${tgt}) endif() diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt new file mode 100644 index 0000000000..99b40e57e7 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt @@ -0,0 +1,5 @@ +if(USE_TPU) + file(GLOB XPORT_SOURCES "*.cpp") + add_library(tent_xport_tpu STATIC ${XPORT_SOURCES}) + target_link_libraries(tent_xport_tpu PUBLIC tent_rpc tent_common) +endif() diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp new file mode 100644 index 0000000000..c3f7646497 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp @@ -0,0 +1,165 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/tpu/tpu_transport.h" + +#include + +#include "tent/common/status.h" +#include "tent/runtime/platform.h" +#include "tent/runtime/slab.h" + +namespace mooncake { +namespace tent { + +TpuTransport::TpuTransport() : installed_(false) {} + +TpuTransport::~TpuTransport() { uninstall(); } + +Status TpuTransport::install(std::string &local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf) { + if (installed_) { + return Status::InvalidArgument( + "TPU transport has been installed" LOC_MARK); + } + metadata_ = metadata; + local_segment_name_ = local_segment_name; + local_topology_ = local_topology; + conf_ = conf; + installed_ = true; + // TPU HBM is not NIC-addressable: only the device<->host staging hop is + // supported here. gpu_to_gpu stays false so the engine always stages + // cross-node traffic through host DRAM. + caps.gpu_to_dram = true; + caps.dram_to_gpu = true; + return Status::OK(); +} + +Status TpuTransport::uninstall() { + if (installed_) { + metadata_.reset(); + installed_ = false; + } + return Status::OK(); +} + +Status TpuTransport::allocateSubBatch(SubBatchRef &batch, size_t max_size) { + auto tpu_batch = Slab::Get().allocate(); + if (!tpu_batch) + return Status::InternalError("Unable to allocate TPU sub-batch"); + batch = tpu_batch; + tpu_batch->task_list.reserve(max_size); + tpu_batch->max_size = max_size; + return Status::OK(); +} + +Status TpuTransport::freeSubBatch(SubBatchRef &batch) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + Slab::Get().deallocate(tpu_batch); + batch = nullptr; + return Status::OK(); +} + +Status TpuTransport::submitTransferTasks( + SubBatchRef batch, const std::vector &request_list) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + if (request_list.size() + tpu_batch->task_list.size() > tpu_batch->max_size) + return Status::TooManyRequests("Exceed batch capacity" LOC_MARK); + for (auto &request : request_list) { + tpu_batch->task_list.push_back(TpuTask{}); + auto &task = tpu_batch->task_list[tpu_batch->task_list.size() - 1]; + task.request = request; + task.status_word = TransferStatusEnum::PENDING; + task.transferred_bytes = 0; + startTransfer(&task, tpu_batch); + } + return Status::OK(); +} + +void TpuTransport::startTransfer(TpuTask *task, TpuSubBatch *batch) { + // TpuTransport only handles the local device<->host staging hop, so the + // target is always the local staging buffer (LOCAL_SEGMENT_ID). Anything + // else indicates a routing bug: TPU HBM cannot be a remote transfer peer. + if (task->request.target_id != LOCAL_SEGMENT_ID) { + LOG(ERROR) << "TpuTransport: unexpected non-local target " + << task->request.target_id + << "; TPU only supports local staging copies"; + task->status_word = TransferStatusEnum::FAILED; + task->transferred_bytes = 0; + batch->notifyProgress(); + return; + } + + void *staging = reinterpret_cast(task->request.target_offset); + Status status; + if (task->request.opcode == Request::READ) + // host staging buffer -> device (H2D) + status = Platform::getLoader().copy(task->request.source, staging, + task->request.length); + else + // device -> host staging buffer (D2H) + status = Platform::getLoader().copy(staging, task->request.source, + task->request.length); + + if (status.ok()) { + task->transferred_bytes = task->request.length; + task->status_word = TransferStatusEnum::COMPLETED; + } else { + LOG(WARNING) << "TpuTransport: staging copy failed: " + << status.ToString(); + task->status_word = TransferStatusEnum::FAILED; + } + batch->notifyProgress(); +} + +Status TpuTransport::getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus &status) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + if (task_id < 0 || task_id >= (int)tpu_batch->task_list.size()) { + return Status::InvalidArgument("Invalid task id" LOC_MARK); + } + auto &task = tpu_batch->task_list[task_id]; + status = TransferStatus{task.status_word, task.transferred_bytes}; + return Status::OK(); +} + +Status TpuTransport::addMemoryBuffer(BufferDesc &desc, + const MemoryOptions &options) { + LocationParser location(desc.location); + // Tag both TPU device buffers and host staging buffers so the staging + // policy can route the local HBM<->host hop through this transport. Routing + // keys on the target (host) buffer's transports, so the host buffer must + // carry TransportType::TPU as well (mirrors NVLinkTransport). + if (location.type() != "tpu" && location.type() != "cpu" && + location.type() != kWildcardLocation) { + return Status::OK(); // Not our buffer; leave it untagged. + } + desc.transports.push_back(TransportType::TPU); + return Status::OK(); +} + +Status TpuTransport::removeMemoryBuffer(BufferDesc &desc) { + return Status::OK(); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 4bd5553035..fa70ae824d 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -190,3 +190,26 @@ target_include_directories(tent_runtime_queue_dispatch_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_runtime_queue_dispatch_test COMMAND tent_runtime_queue_dispatch_test) + +# TPU PJRT shim test: exercises the dlopen'd adapter ABI (device-pointer +# classification, D2H/H2D copy, device topology) against an in-process mock +# adapter, so it runs on any Linux host without TPU hardware or a PJRT runtime. +if(USE_TPU) + # Mock adapter satisfying the tpu_pjrt_abi.h C ABI. Built as a shared object + # that the shim (and the test) load via dlopen. + add_library(mock_tpu_pjrt SHARED tpu/mock_tpu_pjrt_adapter.cpp) + target_include_directories(mock_tpu_pjrt + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + + add_executable(tent_tpu_pjrt_shim_test tpu/tpu_pjrt_shim_test.cpp) + add_dependencies(tent_tpu_pjrt_shim_test mock_tpu_pjrt) + target_link_libraries(tent_tpu_pjrt_shim_test + PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) + target_include_directories(tent_tpu_pjrt_shim_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_compile_definitions( + tent_tpu_pjrt_shim_test + PRIVATE MOCK_TPU_PJRT_LIB="$") + add_test(NAME tent_tpu_pjrt_shim_test COMMAND tent_tpu_pjrt_shim_test) +endif() diff --git a/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp new file mode 100644 index 0000000000..48f569337b --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp @@ -0,0 +1,90 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Mock TPU/PJRT adapter for unit testing TpuPjrtShim and TpuPlatform without +// TPU hardware or a PJRT runtime. It implements the C ABI declared in +// tpu_pjrt_abi.h using ordinary host memory: "device" buffers are plain +// allocations the test registers via the mock_* test helpers below, and D2H/H2D +// copies are memcpy. This lets the routing and shim logic be exercised on any +// Linux host (see tpu_pjrt_shim_test.cpp). + +#include +#include +#include + +#include "tent/platform/tpu_pjrt_abi.h" + +namespace { +std::mutex g_mutex; +// Registered fake device buffers: pointer -> device ordinal. +std::unordered_map g_device_registry; +int g_device_count = 4; +} // namespace + +extern "C" { + +// --- ABI required by TpuPjrtShim ------------------------------------------- + +int mc_tpu_pjrt_init(void) { return 0; } + +int mc_tpu_pjrt_is_device_ptr(const void *addr) { + std::lock_guard lock(g_mutex); + return g_device_registry.count(addr) ? 1 : 0; +} + +int mc_tpu_pjrt_device_index(const void *addr) { + std::lock_guard lock(g_mutex); + auto it = g_device_registry.find(addr); + return it == g_device_registry.end() ? -1 : it->second; +} + +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len) { + if (!host_dst || !device_src) return 1; + std::memcpy(host_dst, device_src, len); + return 0; +} + +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len) { + if (!device_dst || !host_src) return 1; + std::memcpy(device_dst, host_src, len); + return 0; +} + +int mc_tpu_pjrt_device_count(void) { return g_device_count; } + +int mc_tpu_pjrt_device_numa(int index) { + // Deterministic fake affinity: even devices on node 0, odd on node 1. + if (index < 0 || index >= g_device_count) return -1; + return index % 2; +} + +// --- Test-only helpers (not part of the shim ABI) -------------------------- + +void mock_tpu_pjrt_register_device(const void *addr, int index) { + std::lock_guard lock(g_mutex); + g_device_registry[addr] = index; +} + +void mock_tpu_pjrt_reset(void) { + std::lock_guard lock(g_mutex); + g_device_registry.clear(); + g_device_count = 4; +} + +void mock_tpu_pjrt_set_device_count(int count) { + std::lock_guard lock(g_mutex); + g_device_count = count; +} + +} // extern "C" diff --git a/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp new file mode 100644 index 0000000000..e1fa7bc359 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp @@ -0,0 +1,117 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Unit test for TpuPjrtShim against the mock adapter +// (mock_tpu_pjrt_adapter.cpp). Runs on any Linux host: no TPU hardware or PJRT +// runtime required. +// +// MOCK_TPU_PJRT_LIB is the path to the built mock adapter, injected by CMake. + +#include "tent/platform/tpu_pjrt_shim.h" + +#include +#include + +#include +#include + +#ifndef MOCK_TPU_PJRT_LIB +#error "MOCK_TPU_PJRT_LIB must be defined by the build (path to mock adapter)" +#endif + +namespace mooncake { +namespace tent { +namespace { + +using RegisterFn = void (*)(const void *, int); +using ResetFn = void (*)(); +using SetCountFn = void (*)(int); + +// Loads the mock adapter a second time (dlopen is reference-counted, so this is +// the same image the shim loads) to reach the test-only registration helpers. +class TpuShimTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + // Point the shim at the mock adapter before its singleton is first + // used. + ::setenv("MC_TPU_PJRT_LIB", MOCK_TPU_PJRT_LIB, /*overwrite=*/1); + } + + void SetUp() override { + mock_ = dlopen(MOCK_TPU_PJRT_LIB, RTLD_NOW | RTLD_GLOBAL); + ASSERT_NE(mock_, nullptr) << dlerror(); + register_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_register_device")); + reset_ = reinterpret_cast(dlsym(mock_, "mock_tpu_pjrt_reset")); + set_count_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_set_device_count")); + ASSERT_NE(register_, nullptr); + ASSERT_NE(reset_, nullptr); + ASSERT_NE(set_count_, nullptr); + reset_(); + } + + void TearDown() override { + if (reset_) reset_(); + if (mock_) dlclose(mock_); + } + + void *mock_ = nullptr; + RegisterFn register_ = nullptr; + ResetFn reset_ = nullptr; + SetCountFn set_count_ = nullptr; +}; + +TEST_F(TpuShimTest, AdapterLoadsAndReportsDevices) { + auto &shim = TpuPjrtShim::instance(); + ASSERT_TRUE(shim.available()); + EXPECT_EQ(shim.deviceCount(), 4); + EXPECT_EQ(shim.deviceNumaNode(0), 0); + EXPECT_EQ(shim.deviceNumaNode(1), 1); + EXPECT_EQ(shim.deviceNumaNode(99), -1); +} + +TEST_F(TpuShimTest, ClassifiesRegisteredDevicePointers) { + int host_value = 0; + std::vector fake_device(64); + register_(fake_device.data(), /*index=*/2); + + auto &shim = TpuPjrtShim::instance(); + EXPECT_TRUE(shim.isDevicePtr(fake_device.data())); + EXPECT_EQ(shim.deviceIndex(fake_device.data()), 2); + + // Unregistered host memory is not device memory. + EXPECT_FALSE(shim.isDevicePtr(&host_value)); + EXPECT_EQ(shim.deviceIndex(&host_value), -1); + EXPECT_FALSE(shim.isDevicePtr(nullptr)); +} + +TEST_F(TpuShimTest, CopyRoundTripMovesBytes) { + auto &shim = TpuPjrtShim::instance(); + const std::vector src = {1, 2, 3, 4, 5, 6, 7, 8}; + std::vector device(src.size(), 0); + std::vector dst(src.size(), 0); + + // host -> "device" + ASSERT_TRUE(shim.copyH2D(device.data(), src.data(), src.size()).ok()); + EXPECT_EQ(device, src); + + // "device" -> host + ASSERT_TRUE(shim.copyD2H(dst.data(), device.data(), device.size()).ok()); + EXPECT_EQ(dst, src); +} + +} // namespace +} // namespace tent +} // namespace mooncake From 8567ab027f7b6726e3536affce4ca3fce73b203d Mon Sep 17 00:00:00 2001 From: smartssw <130918906+smartssw@users.noreply.github.com> Date: Sat, 4 Jul 2026 23:08:00 +0800 Subject: [PATCH 03/57] [Store] Make batch_query_keys read-only and return all replica types (#2685) --- .../source/http-api-reference/http-service.md | 23 ++- mooncake-store/include/master_service.h | 18 +++ mooncake-store/include/rpc_service.h | 9 ++ mooncake-store/src/master_admin_service.cpp | 134 +++++++++++------- mooncake-store/src/master_service.cpp | 116 +++++++++++++++ mooncake-store/src/rpc_service.cpp | 15 ++ .../tests/master_admin_server_test.cpp | 49 +++++++ .../tests/promotion_on_hit_test.cpp | 85 +++++++++++ 8 files changed, 399 insertions(+), 50 deletions(-) diff --git a/docs/source/http-api-reference/http-service.md b/docs/source/http-api-reference/http-service.md index cbd79e007d..8c8ee90fdd 100644 --- a/docs/source/http-api-reference/http-service.md +++ b/docs/source/http-api-reference/http-service.md @@ -91,7 +91,7 @@ curl "http://localhost:8080/query_key?key=my_object" ``` #### `/batch_query_keys` -Retrieve replica information for multiple keys in a single request, including memory locations and transport endpoints for each key. +Retrieve replica information for multiple keys in a single request, including memory locations and transport endpoints for each key. The endpoint performs a read-only metadata lookup and does not grant leases, trigger promotion, or update cache-hit metrics. **Method**: `GET` **Parameters**: `keys` (query parameter) - Comma-separated list of object keys to query (format: key1,key2,key3) @@ -115,6 +115,25 @@ curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3" "transport_endpoint_": "hostname:port", "buffer_descriptor": {...} } + ], + "disk_values": [ + { + "file_path": "/path/to/object", + "object_size": 4096 + } + ], + "local_disk_values": [ + { + "client_id": "12345-67890", + "object_size": 4096, + "transport_endpoint": "hostname:port" + } + ], + "nof_values": [ + { + "transport_endpoint_": "hostname:port", + "buffer_descriptor": {...} + } ] }, "key2": { @@ -125,6 +144,8 @@ curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3" } ``` +The `values` field is always present (empty array when no memory replica exists). The `disk_values`, `local_disk_values`, and `nof_values` fields are optional and only appear when the corresponding replica type is present for the key. + #### `/get_all_keys` List all keys currently stored in the distributed system. diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 1479c96fdc..b5591d6507 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -322,6 +322,15 @@ class MasterService { auto GetReplicaList(const std::string& key, const std::string& tenant_id) -> tl::expected; + /** + * @brief Read-only single-key replica list query for admin use. + * Unlike GetReplicaList, this does not grant leases, trigger + * promotion, or update cache-hit metrics. + */ + auto GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) + -> tl::expected; + /** * @brief Get replica lists for a batch of objects. */ @@ -329,6 +338,15 @@ class MasterService { BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id); + /** + * @brief Read-only batch replica list query for admin use. + * Unlike BatchGetReplicaList, this does not grant leases, trigger + * promotion, or update cache-hit metrics. + */ + std::vector> + BatchGetReplicaListForAdmin(const std::vector& keys, + const std::string& tenant_id); + /** * @brief Start a put operation for an object * @param[out] replica_list Vector to store replica information for the diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index bf93f386f4..f54f2ec0dd 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -65,6 +65,15 @@ class WrappedMasterService { BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id = "default"); + // Read-only admin variants: no lease grants, no promotion, no metric + // updates. + std::vector> + BatchGetReplicaListForAdmin(const std::vector& keys, + const std::string& tenant_id = "default"); + + tl::expected GetReplicaListForAdmin( + const std::string& key, const std::string& tenant_id = "default"); + tl::expected, ErrorCode> PutStart( const UUID& client_id, const std::string& key, const uint64_t slice_length, const ReplicateConfig& config, diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp index 0898dadb8b..cb221af5de 100644 --- a/mooncake-store/src/master_admin_service.cpp +++ b/mooncake-store/src/master_admin_service.cpp @@ -920,12 +920,30 @@ void MasterAdminServer::HandleSegmentStatus( }); } +struct HttpDiskReplicaInfo { + std::string file_path; + uint64_t object_size = 0; + YLT_REFL(HttpDiskReplicaInfo, file_path, object_size); +}; + +struct HttpLocalDiskReplicaInfo { + std::string client_id; + uint64_t object_size = 0; + std::string transport_endpoint; + YLT_REFL(HttpLocalDiskReplicaInfo, client_id, object_size, + transport_endpoint); +}; + struct HttpBatchQueryKeyResult { bool ok{false}; std::optional error; std::optional> values; + std::optional> disk_values; + std::optional> local_disk_values; + std::optional> nof_values; }; -YLT_REFL(HttpBatchQueryKeyResult, ok, error, values); +YLT_REFL(HttpBatchQueryKeyResult, ok, error, values, disk_values, + local_disk_values, nof_values); struct HttpBatchQueryKeysResponse { bool success{false}; @@ -935,63 +953,81 @@ YLT_REFL(HttpBatchQueryKeysResponse, success, data); void MasterAdminServer::HandleBatchQueryKeys( coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { - auto service = GetActiveService(); - if (!service) { - WriteSimpleErrorResponse(resp, - coro_http::status_type::service_unavailable, - "service plane is not active"); - return; - } + WithActiveService(resp, [&](auto service) { + auto keys_str = req.get_decode_query_value("keys"); + std::vector keys; + if (!keys_str.empty()) { + std::string_view sv(keys_str); + size_t pos = 0; + while ((pos = sv.find(',')) != std::string_view::npos) { + keys.emplace_back(sv.substr(0, pos)); + sv.remove_prefix(pos + 1); + } + keys.emplace_back(sv); + } - auto keys_str = req.get_query_value("keys"); - std::vector keys; - if (!keys_str.empty()) { - std::string_view sv(keys_str); - size_t pos = 0; - while ((pos = sv.find(',')) != std::string_view::npos) { - keys.emplace_back(sv.substr(0, pos)); - sv.remove_prefix(pos + 1); + if (keys.empty()) { + WriteSimpleErrorResponse( + resp, coro_http::status_type::bad_request, + "No keys provided. Use ?keys=key1,key2,..."); + return; } - keys.emplace_back(sv); - } - if (keys.empty()) { - WriteSimpleErrorResponse(resp, coro_http::status_type::bad_request, - "No keys provided. Use ?keys=key1,key2,..."); - return; - } + auto results = service->BatchGetReplicaListForAdmin(keys, "default"); + const size_t n = std::min(keys.size(), results.size()); + HttpBatchQueryKeysResponse payload; + payload.success = true; - auto results = service->BatchGetReplicaList(keys, "default"); - const size_t n = std::min(keys.size(), results.size()); - HttpBatchQueryKeysResponse payload; - payload.success = true; + for (size_t i = 0; i < n; ++i) { + const auto& result = results[i]; + HttpBatchQueryKeyResult item; + if (!result.has_value()) { + item.error = toString(result.error()); + payload.data.emplace(keys[i], std::move(item)); + continue; + } - for (size_t i = 0; i < n; ++i) { - const auto& result = results[i]; - HttpBatchQueryKeyResult item; - if (!result.has_value()) { - item.error = toString(result.error()); + item.ok = true; + item.values = std::vector{}; + for (const auto& replica : result.value().replicas) { + if (replica.is_memory_replica()) { + item.values->emplace_back( + replica.get_memory_descriptor().buffer_descriptor); + } else if (replica.is_disk_replica()) { + if (!item.disk_values.has_value()) { + item.disk_values = std::vector{}; + } + auto& d = replica.get_disk_descriptor(); + item.disk_values->emplace_back( + HttpDiskReplicaInfo{d.file_path, d.object_size}); + } else if (replica.is_local_disk_replica()) { + if (!item.local_disk_values.has_value()) { + item.local_disk_values = + std::vector{}; + } + auto& d = replica.get_local_disk_descriptor(); + item.local_disk_values->emplace_back( + HttpLocalDiskReplicaInfo{UuidToString(d.client_id), + d.object_size, + d.transport_endpoint}); + } else if (replica.is_nof_replica()) { + if (!item.nof_values.has_value()) { + item.nof_values = + std::vector{}; + } + item.nof_values->emplace_back( + replica.get_nof_descriptor().buffer_descriptor); + } + } payload.data.emplace(keys[i], std::move(item)); - continue; } - item.ok = true; - item.values = std::vector{}; - for (const auto& replica : result.value().replicas) { - if (!replica.is_memory_replica()) { - continue; - } - item.values->emplace_back( - replica.get_memory_descriptor().buffer_descriptor); + if (results.size() != keys.size()) { + LOG(WARNING) << "BatchGetReplicaListForAdmin size mismatch: keys=" + << keys.size() << " results=" << results.size(); } - payload.data.emplace(keys[i], std::move(item)); - } - - if (results.size() != keys.size()) { - LOG(WARNING) << "BatchGetReplicaList size mismatch: keys=" - << keys.size() << " results=" << results.size(); - } - WriteJsonResponse(resp, coro_http::status_type::ok, payload); + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); } void MasterAdminServer::HandleGetTenantQuotas( diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index c051fe1682..5010c7053c 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -2567,6 +2567,35 @@ auto MasterService::GetReplicaList(const std::string& key, return resp; } +auto MasterService::GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) + -> tl::expected { + const auto object_id = MakeObjectIdentity(key, tenant_id); + + std::shared_lock shared_lock(snapshot_mutex_); + MetadataAccessorRO accessor(this, object_id); + + if (!accessor.Exists()) { + VLOG(1) << "key=" << key << ", info=object_not_found"; + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + const auto& metadata = accessor.Get(); + + std::vector replica_list; + metadata.VisitReplicas( + &Replica::fn_is_completed, [&replica_list](const Replica& replica) { + replica_list.emplace_back(replica.get_descriptor()); + }); + + if (replica_list.empty()) { + LOG(WARNING) << "key=" << key << ", error=replica_not_ready"; + return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); + } + + return GetReplicaListResponse(std::move(replica_list), + default_kv_lease_ttl_); +} + std::vector> MasterService::BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id) { @@ -2688,6 +2717,93 @@ MasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +MasterService::BatchGetReplicaListForAdmin(const std::vector& keys, + const std::string& tenant_id) { + using GetResult = tl::expected; + + std::vector results( + keys.size(), tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND)); + if (keys.empty()) { + return results; + } + + const auto normalized_tenant = NormalizeTenantId(tenant_id); + constexpr size_t kInvalidKeyIndex = std::numeric_limits::max(); + std::array key_list_heads; + key_list_heads.fill(kInvalidKeyIndex); + std::vector next_key_indexes(keys.size(), kInvalidKeyIndex); + { + std::shared_lock lock(group_routing_mutex_); + for (size_t i = keys.size(); i > 0; --i) { + const size_t original_idx = i - 1; + const auto scoped_key = + MakeTenantScopedKey(normalized_tenant, keys[original_idx]); + const auto route_it = object_group_ids_.find(scoped_key); + const size_t shard_idx = + route_it == object_group_ids_.end() + ? getShardIndex(normalized_tenant, keys[original_idx]) + : getShardIndex(route_it->second); + next_key_indexes[original_idx] = key_list_heads[shard_idx]; + key_list_heads[shard_idx] = original_idx; + } + } + + const size_t start_shard = RandomIndex(kNumShards); + for (size_t scanned = 0; scanned < kNumShards; ++scanned) { + const size_t shard_idx = + (start_shard + kNumShards - scanned) % kNumShards; + if (key_list_heads[shard_idx] == kInvalidKeyIndex) { + continue; + } + + std::shared_lock shared_lock(snapshot_mutex_); + { + MetadataShardAccessorRO shard(this, shard_idx); + const auto tenant_it = shard->tenants.find(normalized_tenant); + for (size_t original_idx = key_list_heads[shard_idx]; + original_idx != kInvalidKeyIndex; + original_idx = next_key_indexes[original_idx]) { + const std::string& key = keys[original_idx]; + + if (tenant_it == shard->tenants.end()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; + } + + const auto& tenant_state = tenant_it->second; + const auto metadata_it = tenant_state.metadata.find(key); + if (metadata_it == tenant_state.metadata.end() || + !metadata_it->second.IsValid()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; + } + + const auto& metadata = metadata_it->second; + std::vector replica_list; + metadata.VisitReplicas( + &Replica::fn_is_completed, + [&replica_list](const Replica& replica) { + replica_list.emplace_back(replica.get_descriptor()); + }); + + if (replica_list.empty()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); + continue; + } + + results[original_idx] = GetReplicaListResponse( + std::move(replica_list), default_kv_lease_ttl_); + } + } + } + + return results; +} + auto MasterService::AllocateAndInsertMetadata( MetadataShardAccessorRW& shard, const UUID& client_id, const std::string& key, uint64_t value_length, diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 493a496c23..5d2ae9efad 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -230,6 +230,21 @@ WrappedMasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +WrappedMasterService::BatchGetReplicaListForAdmin( + const std::vector& keys, const std::string& tenant_id) { + return master_service_.BatchGetReplicaListForAdmin(keys, tenant_id); +} + +tl::expected +WrappedMasterService::GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) { + return execute_rpc( + "GetReplicaListForAdmin", + [&] { return master_service_.GetReplicaListForAdmin(key, tenant_id); }, + [&](auto& timer) { timer.LogRequest("key=", key); }, [] {}, [] {}); +} + tl::expected, ErrorCode> WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, const uint64_t slice_length, diff --git a/mooncake-store/tests/master_admin_server_test.cpp b/mooncake-store/tests/master_admin_server_test.cpp index 5ef0cd62a7..20249277cd 100644 --- a/mooncake-store/tests/master_admin_server_test.cpp +++ b/mooncake-store/tests/master_admin_server_test.cpp @@ -1206,6 +1206,55 @@ TEST_F(MasterAdminServerTest, MultipleSegmentsAndKeys) { admin.Stop(); } +// /batch_query_keys returns replica metadata for disk-based keys via the +// optional disk_values/local_disk_values/nof_values fields, while the existing +// values field stays present (empty array) for backward compatibility. +TEST_F(MasterAdminServerTest, BatchQueryKeysReturnsLocalDiskReplicaInfo) { + WrappedMasterServiceConfig svc_config; + svc_config.default_kv_lease_ttl = 5000; + svc_config.enable_metric_reporting = false; + svc_config.enable_offload = true; // required to mount local-disk segments + auto service = std::make_shared(svc_config); + + UUID client_id = generate_uuid(); + Segment segment; + segment.id = generate_uuid(); + segment.name = "ld_segment"; + segment.base = 0x600000000; + segment.size = 8 * 1024 * 1024; + ASSERT_TRUE(service->MountSegment(segment, client_id).has_value()); + ASSERT_TRUE(service->MountLocalDiskSegment(client_id, true).has_value()); + + const std::string key = "ld_only_key"; + const std::string endpoint = "127.0.0.1:9999"; + StorageObjectMetadata sm; + sm.bucket_id = 0; + sm.offset = 0; + sm.key_size = static_cast(key.size()); + sm.data_size = 2048; + sm.transport_endpoint = endpoint; + OffloadTaskItem task{.tenant_id = "default", .key = key, .size = 2048}; + ASSERT_TRUE( + service->NotifyOffloadSuccess(client_id, {task}, {sm}).has_value()); + + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + admin.SetServiceDelegate(service); + admin.SetServiceAvailable(true); + + auto resp = HttpGet(port, "/batch_query_keys?keys=" + key); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("\"success\":true"), std::string::npos); + EXPECT_NE(resp.body.find("local_disk_values"), std::string::npos); + EXPECT_NE(resp.body.find(endpoint), std::string::npos); + // Backward compat: a disk-only key still reports an (empty) values array. + EXPECT_NE(resp.body.find("\"values\":[]"), std::string::npos); + + admin.Stop(); +} + } // namespace test } // namespace mooncake diff --git a/mooncake-store/tests/promotion_on_hit_test.cpp b/mooncake-store/tests/promotion_on_hit_test.cpp index 2743c4e252..bc2840970c 100644 --- a/mooncake-store/tests/promotion_on_hit_test.cpp +++ b/mooncake-store/tests/promotion_on_hit_test.cpp @@ -283,6 +283,91 @@ TEST_F(PromotionOnHitTest, BatchGetReplicaListPromotesLocalDiskOnlyObject) { service->RemoveAll(); } +// The read-only admin batch query must NOT trigger promotion-on-hit, even for a +// LOCAL_DISK-only key. Contrast with BatchGetReplicaListPromotesLocalDiskOnly +// Object above, where the client-facing BatchGetReplicaList admits the key for +// promotion. +TEST_F(PromotionOnHitTest, + BatchGetReplicaListForAdminDoesNotPromoteLocalDiskOnlyObject) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.promotion_max_per_heartbeat = 2; + config.default_kv_lease_ttl = 2000; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto ctx = PrepareSegment(*service, "test_segment_admin", + kDefaultSegmentBase, seg_size); + + const std::string key = "k_batch_admin_no_promote"; + ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, key, 1024, + ctx.segment_name)); + + auto& mm = MasterMetricManager::instance(); + const int64_t admitted_pre = mm.get_promotion_admitted(); + const int64_t in_flight_pre = mm.get_promotion_in_flight(); + + auto result = service->BatchGetReplicaListForAdmin( + std::vector{key}, "default"); + ASSERT_EQ(result.size(), 1u); + ASSERT_TRUE(result[0].has_value()); + ASSERT_EQ(result[0]->replicas.size(), 1u); + EXPECT_TRUE(result[0]->replicas[0].is_local_disk_replica()); + + // No promotion admitted or enqueued by the read-only admin path. + EXPECT_EQ(mm.get_promotion_admitted() - admitted_pre, 0); + EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 0); + auto pending = service->PromotionObjectHeartbeat(ctx.client_id); + ASSERT_TRUE(pending.has_value()); + EXPECT_EQ(pending->size(), 0u); + EXPECT_EQ(CountPromotionTask(*pending, key), 0u); + + service->RemoveAll(); +} + +// The read-only admin batch query must NOT update the store-observed cache-hit +// counters. Contrast with the client-facing BatchGetReplicaList, which bumps +// the memory-cache-hit counter for a MEMORY replica. +TEST_F(PromotionOnHitTest, + BatchGetReplicaListForAdminDoesNotUpdateCacheHitMetrics) { + MasterServiceConfig config; + config.enable_offload = true; + config.default_kv_lease_ttl = 2000; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto ctx = PrepareSegment(*service, "metrics_segment", kDefaultSegmentBase, + seg_size); + + const std::string key = "k_admin_no_metric"; + PutObject(*service, ctx.client_id, key, 1024); + + using CacheHitStat = MasterMetricManager::CacheHitStat; + auto mem_hits = []() { + auto stats = MasterMetricManager::instance().calculate_cache_stats(); + return stats[CacheHitStat::MEMORY_HITS]; + }; + + const double before_admin = mem_hits(); + + // Read-only admin query must leave the memory-cache-hit counter untouched. + auto admin_result = service->BatchGetReplicaListForAdmin( + std::vector{key}, "default"); + ASSERT_EQ(admin_result.size(), 1u); + ASSERT_TRUE(admin_result[0].has_value()); + EXPECT_EQ(mem_hits(), before_admin); + + // The client-facing path does bump it, proving the assertion above is + // meaningful rather than a counter that never moves. + (void)service->BatchGetReplicaList(std::vector{key}, + "default"); + EXPECT_GT(mem_hits(), before_admin); + + service->RemoveAll(); +} + // PromotionObjectHeartbeat returns an empty task list when called against a // client that has no LocalDiskSegment registered. TEST_F(PromotionOnHitTest, HeartbeatReturnsErrorForUnknownClient) { From f704fb83fe229781027d4bb36b575414bcb623c5 Mon Sep 17 00:00:00 2001 From: Aoi Date: Sat, 4 Jul 2026 23:15:16 +0800 Subject: [PATCH 04/57] [Doc] Add agent guidance (#2712) --- .gitignore | 3 -- AGENTS.md | 25 ++++++++++++++ CLAUDE.md | 1 + CONTRIBUTING.md | 8 ++++- docs/AGENTS.md | 90 +++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 AGENTS.md create mode 100644 CLAUDE.md create mode 100644 docs/AGENTS.md diff --git a/.gitignore b/.gitignore index 4a3ec467a5..35a253a732 100644 --- a/.gitignore +++ b/.gitignore @@ -198,9 +198,6 @@ mooncake-wheel/mooncake/allocator_ascend_npu.py mooncake-wheel/mooncake/mooncake_master mooncake-wheel/mooncake/transfer_engine_bench -# Claude Code Memory -CLAUDE.md - # CodeQL _codeql_detected_source_root diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..f7332c2d0a --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,25 @@ +# AGENTS.md + +## `docs/` Directory Changes + +- Before modifying files under `docs/`, read `docs/AGENTS.md`. + +## Pull Request Guidelines + +- Follow `CONTRIBUTING.md` for PR title prefixes, RFC expectations, and + contribution workflow. +- Before opening a PR for nontrivial work, check whether an existing issue or + open PR already covers the same change. If the work overlaps, explain the + difference instead of duplicating it. +- Do not open low-value busywork PRs for isolated typo, style, or mechanical + changes unless they are part of a substantive requested change. +- Use `.github/pull_request_template.md` when preparing a PR, and fill in the + relevant sections for description, module, type of change, testing, + checklist, and AI assistance disclosure. +- For AI-assisted changes, make sure the human submitter has reviewed every + changed line and can defend the change end-to-end. +- Run pre-commit locally on the files touched by the change before handoff when + the toolchain is available. If broader hooks or `pre-commit run --all-files` + rewrite unrelated files, do not include those unrelated edits in the PR. +- Keep PRs lean: review `git diff` before staging, and include only changes + required for the requested task. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..43c994c2d3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5ce88ac0f3..4116d4bf02 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -16,7 +16,8 @@ Thank you for your interest in contributing to Mooncake! Our community warmly we ### PR Title and Classification -Use a prefixed PR title to indicate the type of changes. Please use one of the following: +Use a prefixed PR title to indicate the type or module affected by the changes. +Prefer one of the following documented prefixes: - ``[Bugfix]`` for bug fixes. - ``[CI/Build]`` for build or continuous integration improvements. @@ -28,6 +29,11 @@ Use a prefixed PR title to indicate the type of changes. Please use one of the f - ``[Misc]`` for PRs that do not fit the above categories. Please use this sparingly. +The project history also contains common aliases and module prefixes. Use these +when they better match the change scope: ``[Bug fix]``, ``[Build]``, ``[CI]``, +``[Docs]``, ``[EP]``, ``[Feature]``, ``[MUSA]``, ``[PG]``, ``[TE]``, +``[TENT]``, and ``[Wheel]``. + ### RFC Discussion For major architectural changes (>500 LOC excluding tests), we would expect a GitHub issue (RFC) discussing the technical design and justification. diff --git a/docs/AGENTS.md b/docs/AGENTS.md new file mode 100644 index 0000000000..debb165ca2 --- /dev/null +++ b/docs/AGENTS.md @@ -0,0 +1,90 @@ +# AGENTS.md - Mooncake Documentation + +This file gives coding agents the repo-local rules for modifying files under +`docs/`. Keep `README.md` as the human-facing quickstart. Use this file for +agent workflow, verification, and maintenance guidance. + +## Scope + +- Applies to changes under `docs/`, especially `docs/source/`. +- Prefer small, reviewable documentation changes. +- Do not rewrite unrelated pages, generated files, or formatting-only content. +- Preserve existing documentation structure unless the user asks for a broader + reorganization. + +## Build and Preview + +Run documentation commands from the `docs` directory: + +``` +cd docs +``` + +Install dependencies with `uv` when needed. The requirements file is +`docs/requirements-docs.txt`; after `cd docs`, use the local filename. If there +is an existing venv, prefer using the existing one. Otherwise, create one before +installing dependencies: + +``` +uv venv +uv pip install -r requirements-docs.txt +``` + +Clean stale build output when validating navigation, generated API pages, or +theme behavior: + +``` +make clean +``` + +Build HTML before handing off user-visible documentation changes: + +``` +make html +``` + +Set `locale` correctly before building when the change depends on localized +content or translated output. + +Serve the generated site for review: + +``` +python -m http.server -d build/html/ +``` + +The default URL is `http://localhost:8000`. If port 8000 is busy, choose another +available port. + +## Editing Guidance + +- Source pages live under `docs/source/`. +- Keep links relative and Sphinx-compatible unless an external URL is required. +- For navigation changes, inspect `docs/source/index.md` and the relevant + toctree before editing individual pages. +- Use Sphinx-native structure for documentation behavior. Do not use client-side + JavaScript or post-render DOM patches for navigation or theme behavior. +- If the requested behavior is not supported by Sphinx or the active theme, + prefer adding a Sphinx extension/plugin instead of patching rendered HTML. +- When a page should be linked from content but excluded from the main sidebar, + use a content link plus appropriate Sphinx metadata such as `orphan: true` + instead of hiding rendered sidebar nodes. +- Keep homepage toctree depth conservative. Do not increase `index.md` maxdepth + unless the user explicitly asks for deeper landing-page nesting. + +## Validation Checklist + +- Run `make html` for docs changes that affect rendered pages, navigation, + cross-references, or Sphinx configuration. +- Check the generated HTML for the changed pages. +- For sidebar or toctree changes, verify both the article body and left sidebar + render the intended entries. +- If a local preview server is useful for review, start one from `docs/` with + `python -m http.server -d build/html/` or an alternate port. + +## Pull Request Hygiene + +- Keep docs-only changes narrowly scoped. +- Review `git diff` before staging so generated files or hook-only formatting + changes do not leak into the PR. +- If opening a PR, use the repository pull request template. +- Use the repository PR title prefix rules from the root `AGENTS.md`. From 275e84cf7e4ac357c343cc1eabd8d2f1799aa6b0 Mon Sep 17 00:00:00 2001 From: Zoee <30841158+n-WN@users.noreply.github.com> Date: Sun, 5 Jul 2026 11:49:26 +0800 Subject: [PATCH 05/57] [TransferEngine] Prefer private-range IPv4 GIDs over link-local IPv6 in auto-selection (#2741) * [TransferEngine] Prefer private-range IPv4 GIDs over link-local IPv6 in auto-selection Fixes #2729. isOverlayIPv4() classifies every RFC1918/CGNAT address as overlay, which dropped routable 10.x datacenter-fabric GIDs into the same degraded tier as link-local fe80:: GIDs; the lowest-gid-index tie-break then deterministically picked the link-local entry (indices 0/1 on mlx5), which can only ever work same-L2 and broke RoCEv2 deployments addressed from 10/8. Split the degraded tier instead of reordering anything else: private-range IPv4-mapped GIDs (10/8, 172.16/12, 100.64/10) now rank in their own tier strictly below genuinely routable GIDs and strictly above link-local / overlay-named-interface GIDs. The interface-name overlay heuristic (docker*/cni*/...) remains the strongest demotion signal. All pre-existing cross-tier orderings are preserved; deployments that pin MC_GID_INDEX are unaffected. For the topology reported in #2729 this restores the effective pre-4d7c1a19 selection. --- .../transport/rdma_transport/rdma_gid_probe.h | 53 ++++-- .../transport/rdma_transport/rdma_context.cpp | 2 + .../tests/rdma_gid_probe_test.cpp | 159 ++++++++++++++++++ 3 files changed, 199 insertions(+), 15 deletions(-) diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h index a4ee60dbd3..1927998ade 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h @@ -29,9 +29,18 @@ namespace mooncake { enum class AutoGidCandidateClass { kNetworkRoutable = 0, kNoNetworkRoutable = 1, - kNetworkDegraded = 2, - kNoNetworkDegraded = 3, - kFallbackNonzero = 4, + // Private-range (RFC 1918 / CGNAT) IPv4-mapped GIDs. Split out of the + // degraded tier (#2729): datacenter RoCE fabrics routinely address NICs + // from 10/8, and such a GID is at worst same-subnet-usable — unlike a + // link-local fe80:: GID, which can never work across subnets. The + // relative order of all pre-existing tiers is unchanged; only the former + // tie between private-range IPv4 and link-local/overlay entries within + // "degraded" is split. + kNetworkPrivateV4 = 2, + kNetworkDegraded = 3, + kNoNetworkPrivateV4 = 4, + kNoNetworkDegraded = 5, + kFallbackNonzero = 6, }; enum class AutoGidRetryAction { @@ -72,8 +81,12 @@ inline const char* autoGidCandidateClassToString( return "network-routable"; case AutoGidCandidateClass::kNoNetworkRoutable: return "no-network-routable"; + case AutoGidCandidateClass::kNetworkPrivateV4: + return "network-private-v4"; case AutoGidCandidateClass::kNetworkDegraded: return "network-degraded"; + case AutoGidCandidateClass::kNoNetworkPrivateV4: + return "no-network-private-v4"; case AutoGidCandidateClass::kNoNetworkDegraded: return "no-network-degraded"; case AutoGidCandidateClass::kFallbackNonzero: @@ -95,20 +108,26 @@ inline std::optional classifyAutoGidCandidate( } const bool is_roce_v2 = candidate.gid_type == IBV_GID_TYPE_ROCE_V2; - const bool is_overlay = - is_roce_v2 && (candidate.is_overlay_network || - (candidate.is_ipv4_mapped && candidate.is_overlay_ipv4)); + // Interface-name-based overlay detection (docker*/cni*/...) is the + // reliable demotion signal; a private address range alone is not — DC + // RoCE fabrics commonly use 10/8 (#2729). + const bool is_overlay_iface = is_roce_v2 && candidate.is_overlay_network; + const bool is_private_v4 = is_roce_v2 && !is_overlay_iface && + candidate.is_ipv4_mapped && + candidate.is_overlay_ipv4; const bool is_link_local = is_roce_v2 && !candidate.is_ipv4_mapped && candidate.is_link_local_ipv6; - const bool is_degraded = is_overlay || is_link_local; + const bool is_degraded = is_overlay_iface || is_link_local; if (candidate.has_network_device) { - return is_degraded ? AutoGidCandidateClass::kNetworkDegraded - : AutoGidCandidateClass::kNetworkRoutable; + if (is_degraded) return AutoGidCandidateClass::kNetworkDegraded; + if (is_private_v4) return AutoGidCandidateClass::kNetworkPrivateV4; + return AutoGidCandidateClass::kNetworkRoutable; } - return is_degraded ? AutoGidCandidateClass::kNoNetworkDegraded - : AutoGidCandidateClass::kNoNetworkRoutable; + if (is_degraded) return AutoGidCandidateClass::kNoNetworkDegraded; + if (is_private_v4) return AutoGidCandidateClass::kNoNetworkPrivateV4; + return AutoGidCandidateClass::kNoNetworkRoutable; } inline int autoGidCandidateClassPriority( @@ -118,14 +137,18 @@ inline int autoGidCandidateClassPriority( return 0; case AutoGidCandidateClass::kNoNetworkRoutable: return 1; - case AutoGidCandidateClass::kNetworkDegraded: + case AutoGidCandidateClass::kNetworkPrivateV4: return 2; - case AutoGidCandidateClass::kNoNetworkDegraded: + case AutoGidCandidateClass::kNetworkDegraded: return 3; - case AutoGidCandidateClass::kFallbackNonzero: + case AutoGidCandidateClass::kNoNetworkPrivateV4: return 4; + case AutoGidCandidateClass::kNoNetworkDegraded: + return 5; + case AutoGidCandidateClass::kFallbackNonzero: + return 6; } - return 5; + return 7; } inline std::vector rankAutoGidCandidates( diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp index 71fdcc9358..b208b69aef 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -745,9 +745,11 @@ static GidNetworkState autoGidStateFromSelection( const AutoGidSelection &selection) { switch (selection.candidate_class) { case AutoGidCandidateClass::kNetworkRoutable: + case AutoGidCandidateClass::kNetworkPrivateV4: case AutoGidCandidateClass::kNetworkDegraded: return GidNetworkState::GID_WITH_NETWORK; case AutoGidCandidateClass::kNoNetworkRoutable: + case AutoGidCandidateClass::kNoNetworkPrivateV4: case AutoGidCandidateClass::kNoNetworkDegraded: case AutoGidCandidateClass::kFallbackNonzero: return GidNetworkState::GID_WITHOUT_NETWORK; diff --git a/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp index 740222bcd4..0cbb75b819 100644 --- a/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp @@ -14,6 +14,8 @@ #include +#include +#include #include #include "transport/rdma_transport/rdma_gid_probe.h" @@ -399,4 +401,161 @@ TEST(RdmaGidProbeTest, RetryActionRequiresObservedOrReprobedChange) { AutoGidRetryAction::kRetryWithObservedChange); } +// Regression tests for #2729: a routable-fabric private-range IPv4 GID must +// outrank a link-local IPv6 GID instead of tying with it in the degraded +// tier (where the lowest-index tie-break used to pick fe80::). + +// Exactly the GID table from the #2729 report: fe80 v1/v2 at indices 0/1, +// 10.14.x-mapped v1/v2 at indices 2/3, all on the same netdev. RoCE v1 +// entries are filtered by type; index 3 (private v4, RoCE v2) must win over +// index 1 (link-local, RoCE v2). +TEST(RdmaGidProbeTest, PrefersPrivateRangeV4OverLinkLocal) { + std::vector candidates = { + makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V1, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V1, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 3); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkPrivateV4); +} + +// A genuinely routable (non-private) v4 GID still outranks a private-range +// one: the new tier sits strictly between routable and degraded. +TEST(RdmaGidProbeTest, RoutableV4StillOutranksPrivateRangeV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/5, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 5); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkRoutable); +} + +// An overlay-NAMED interface (docker0/cni/...) stays demoted below +// private-range v4 even when its address is v4-mapped: the interface-name +// heuristic remains the strongest demotion signal. +TEST(RdmaGidProbeTest, OverlayInterfaceStaysDemotedBelowPrivateRangeV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/true, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/4, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 4); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkPrivateV4); +} + +// Cross-tier order is preserved from before the split: a link-local GID +// with a netdev still outranks a private-range v4 GID without one, exactly +// as network-degraded outranked no-network-degraded before. +TEST(RdmaGidProbeTest, NetworkLinkLocalStillOutranksNoNetworkPrivateV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/false, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 1); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkDegraded); +} + +// Completeness pin for the tier tables: every class has a distinct +// priority consistent with its enum order and a unique display name, so a +// future class insertion cannot silently create a tie or reuse a label. +TEST(RdmaGidProbeTest, ClassPriorityAndNameTablesAreComplete) { + const AutoGidCandidateClass all[] = { + AutoGidCandidateClass::kNetworkRoutable, + AutoGidCandidateClass::kNoNetworkRoutable, + AutoGidCandidateClass::kNetworkPrivateV4, + AutoGidCandidateClass::kNetworkDegraded, + AutoGidCandidateClass::kNoNetworkPrivateV4, + AutoGidCandidateClass::kNoNetworkDegraded, + AutoGidCandidateClass::kFallbackNonzero, + }; + int expected_priority = 0; + std::set names; + for (auto cls : all) { + EXPECT_EQ(autoGidCandidateClassPriority(cls), expected_priority++); + std::string name = autoGidCandidateClassToString(cls); + EXPECT_NE(name, "unknown"); + EXPECT_TRUE(names.insert(name).second) + << "duplicate class name: " << name; + } +} + +// When only link-local candidates exist, behavior is unchanged: lowest +// index wins within the tier. +TEST(RdmaGidProbeTest, LinkLocalOnlyKeepsLowestIndexTieBreak) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 1); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkDegraded); +} + } // namespace From 65edc678fda35b86fb49d617da717d63ec5b2fc1 Mon Sep 17 00:00:00 2001 From: Xiao You Date: Sun, 5 Jul 2026 11:51:05 +0800 Subject: [PATCH 06/57] [TE] Downgrade unrecognized mem addr log from ERROR to INFO (#2734) When aclrtPointerGetAttributes returns an unknown location type, the transport already falls back to host memory; INFO is sufficient and avoids noisy error logs during normal operation. Co-authored-by: lbjyx --- .../ascend_direct_transport/ascend_direct_transport.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp index f5869b34e1..51e97eeeb5 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp @@ -343,8 +343,8 @@ int AscendDirectTransport::registerLocalMemory(void *addr, size_t length, } else if (attributes.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) { mem_type = adxl::MEM_DEVICE; } else { - LOG(ERROR) << "mem addr:" << addr - << " can not be recognized, try set to host mem."; + LOG(INFO) << "mem addr:" << addr + << " can not be recognized, try set to host mem."; mem_type = adxl::MEM_HOST; } } else { From bb1b0843868c44fdd692c2e589a76d1e447a5c2a Mon Sep 17 00:00:00 2001 From: mikegguo Date: Sun, 5 Jul 2026 13:41:37 +0800 Subject: [PATCH 07/57] [Store] Enable hugepage mmap in allocate_buffer_numa_segments (#2417) Co-authored-by: fatSheep --- mooncake-store/src/utils.cpp | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 7dbaa4b0d2..83fb8a3aa7 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -337,12 +337,21 @@ void *allocate_buffer_numa_segments(size_t total_size, size_t region_size = align_up(total_size / n, page_size); size_t map_size = region_size * n; - // reserve contiguous VMA, no physical pages yet - void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, - MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + // reserve contiguous VMA; use hugepages if page_size indicates so + unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS; + if (page_size == SZ_2MB) { + flags |= MAP_HUGETLB | MAP_HUGE_2MB; + } else if (page_size == SZ_1GB) { + flags |= MAP_HUGETLB | MAP_HUGE_1GB; + } else if (page_size != static_cast(getpagesize())) { + flags |= MAP_HUGETLB; + } + void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, flags, -1, 0); if (ptr == MAP_FAILED) { - LOG(ERROR) << "mmap failed, size=" << map_size << ", errno=" << errno - << " (" << strerror(errno) << ")"; + LOG(ERROR) << "mmap failed (hugepage=" + << ((flags & MAP_HUGETLB) ? "yes" : "no") + << "), size=" << map_size << ", errno=" << errno << " (" + << strerror(errno) << ")"; return nullptr; } From f66d16d76b4041659568e94501d45d58b503513a Mon Sep 17 00:00:00 2001 From: Zoee <30841158+n-WN@users.noreply.github.com> Date: Sun, 5 Jul 2026 15:57:00 +0800 Subject: [PATCH 08/57] [TENT] Fix data race on local SegmentDesc via copy-on-write snapshots (#2714) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Fix data race on local SegmentDesc via copy-on-write snapshots Writers (SegmentTracker add/remove, transport install paths) mutated the local SegmentDesc in place while readers (GetSegmentDesc JSON dump, findBuffer in transports and route resolution) accessed it without any locking — undefined behavior under concurrent register/unregisterLocalMemory (issue #2477). The local desc is now published as immutable copy-on-write snapshots: - SegmentManager::updateLocal() serializes all mutations behind a writer mutex, applies them to a private clone and atomically publishes it. - getLocal() returns a consistent snapshot; readers may observe a stale snapshot, never a torn one. Staleness matches the existing metadata semantics (remote descs are cached with TTL + best-effort invalidation). - Sites that retain BufferDesc* past their lookup scope now pin the snapshot (RouteHint::pin, withCachedSegment pin overload). - The serialized-JSON cache is version-tagged so a slow dump of an old snapshot can never overwrite the cache entry of a newer publication. - SegmentTracker::query() (unused, returned raw pointers into the buffer vector) is removed. Adds segment_manager_test with a writers-vs-readers regression test that fails under the previous in-place mutation (TSAN data race + torn-entry invariant violations). * [TENT] Address review: release/acquire on version counter, addInBatch rollback - Publish local_desc_version_ with memory_order_release (pairs with the acquire load in getLocal): per [atomics.order] a relaxed bump creates no synchronizes-with edge, so on weakly-ordered targets a reader could tag the old snapshot with the new version and its thread-local cache would serve the stale snapshot until the next publication. Same reasoning for the getLocalDumpedJson tag sample. - Roll back duplicate ref_count bumps when the addInBatch registration callback fails, so a failed registration no longer pins buffers forever (pre-existing behavior on main). The bump itself stays *before* the callback: holding ref_count >= 2 under the writer mutex is what prevents a concurrent unregister from erasing the entry between the duplicate check and the commit. - Regression test for the rollback path. * [TENT] Rebase adaptation: pin snapshots in #2691's batch-submit paths #2691 (merged after this branch was cut) added two segment_manager.getLocal().get() sites in the NVLink/MNNVL batch submit paths. Under the copy-on-write snapshot semantics introduced here, that raw pointer is only guaranteed valid while an owning SegmentDescRef is held; hold the reference for the duration of the batch validation loop. --- .../include/tent/runtime/segment_manager.h | 109 +++++- .../include/tent/runtime/segment_tracker.h | 22 +- .../include/tent/transport/rdma/workers.h | 3 + .../tent/src/runtime/control_plane.cpp | 4 +- .../tent/src/runtime/segment_manager.cpp | 84 ++++- .../tent/src/runtime/segment_tracker.cpp | 217 +++++++----- .../tent/src/runtime/transfer_engine_impl.cpp | 42 ++- .../ascend/ascend_direct_transport.cpp | 9 +- .../src/transport/mnnvl/mnnvl_transport.cpp | 9 +- .../src/transport/nvlink/nvlink_transport.cpp | 9 +- .../src/transport/rdma/rdma_transport.cpp | 43 ++- .../tent/src/transport/rdma/workers.cpp | 6 +- .../tent/src/transport/shm/shm_transport.cpp | 6 +- .../tent/tests/CMakeLists.txt | 10 + .../tent/tests/segment_manager_test.cpp | 309 ++++++++++++++++++ 15 files changed, 713 insertions(+), 169 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/segment_manager_test.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h index 98be76bcf4..26fccd234b 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -58,6 +59,11 @@ class SegmentManager { // cache invalidation and retry on stale segment cache. Status getRemoteCached(SegmentDesc *&desc, SegmentID handle); + // Owning-reference variant: the returned SegmentDescRef keeps the desc + // (and raw pointers into it) alive independently of the thread-local + // cache's lifetime. + Status getRemoteCached(SegmentDescRef &desc, SegmentID handle); + Status getRemote(SegmentDescRef &desc, const std::string &segment_name); // Invalidates the thread-local cache. @@ -77,18 +83,33 @@ class SegmentManager { // invalidated, the segment is refetched and the operation is retried. template Status withCachedSegment(SegmentID segment_id, Func operation) { + SegmentDescRef pin; + return withCachedSegment(segment_id, pin, operation); + } + + // Same as above, but additionally hands the caller an owning reference + // to the segment snapshot in `pin`. Use this variant whenever a raw + // pointer obtained inside `operation` (e.g. a BufferDesc* from + // findBuffer()) must remain valid after `operation` returns: the pointer + // is guaranteed valid only for as long as `pin` is held. + template + Status withCachedSegment(SegmentID segment_id, SegmentDescRef &pin, + Func operation) { static_assert( std::is_same_v, Status>, "operation must return Status"); - // Local segment: no cache lookup or retry required. + // Local segment: no cache lookup or retry required. Pin the current + // snapshot so pointers into it survive concurrent updateLocal(). if (segment_id == LOCAL_SEGMENT_ID) { - return operation(getLocal().get()); + pin = getLocal(); + return operation(pin.get()); } // First get a cached version. SegmentDesc *desc = nullptr; - CHECK_STATUS(getRemoteCached(desc, segment_id)); + CHECK_STATUS(getRemoteCached(pin, segment_id)); + desc = pin.get(); // Do operation under the cached segment. Status res = operation(desc); @@ -98,7 +119,8 @@ class SegmentManager { // If result status is IsNeedsRefreshCache, invalidate cache and retry invalidateRemote(segment_id); - CHECK_STATUS(getRemoteCached(desc, segment_id)); + CHECK_STATUS(getRemoteCached(pin, segment_id)); + desc = pin.get(); // Do operation again res = operation(desc); @@ -114,14 +136,60 @@ class SegmentManager { } public: - SegmentDescRef getLocal() { return local_desc_; } + // Returns the current immutable snapshot of the local SegmentDesc. + // + // Snapshots are copy-on-write: once published they are never mutated, so + // the returned SegmentDescRef (and raw pointers into it, e.g. from + // findBuffer()) stays valid and self-consistent for as long as the caller + // holds the reference. A reader may observe a snapshot that is stale with + // respect to a concurrent register/unregister, never a torn one — + // staleness is already a designed-in property of segment metadata (remote + // peers cache descs with a TTL plus best-effort invalidation pushes). + // + // Do NOT write through the returned pointer; all mutations must go + // through updateLocal(). + SegmentDescRef getLocal() { + // Per-thread snapshot cache with version-based invalidation — the + // same pattern getRemoteCached() uses for remote descs. The fast + // path is one relaxed-acquire load; the shared_mutex is only taken + // after a publication. Keyed by a monotonic manager id (not `this`) + // so a recycled allocation can never satisfy a stale cache entry. + struct Cache { + uint64_t manager_id = UINT64_MAX; + uint64_t version = 0; + SegmentDescRef ref; + }; + thread_local Cache cache; + auto version = local_desc_version_.load(std::memory_order_acquire); + if (cache.manager_id != manager_id_ || cache.version != version || + !cache.ref) { + std::shared_lock guard(local_desc_lock_); + cache.ref = local_desc_; + cache.manager_id = manager_id_; + // Tag with the pre-lock version: if a publication raced in + // between, the tag mismatches on the next call and we simply + // refresh again — the cache can serve a fresher snapshot than + // its tag, never a staler one. + cache.version = version; + } + return cache.ref; + } + + // Applies `mutator` to a private clone of the local SegmentDesc and + // atomically publishes the result as the new snapshot. Mutations are + // serialized by an internal writer mutex. If `mutator` returns a non-OK + // status, nothing is published. + // + // This is the only way to modify the local SegmentDesc; see getLocal() + // for the snapshot semantics it guarantees. + Status updateLocal(const std::function &mutator); // Returns a serialized JSON snapshot of local_desc_. The result is cached // and shared across concurrent GetSegmentDesc RPC handlers; the cache is - // invalidated by synchronizeLocal() (which is called on every register / - // unregisterLocalMemory). This avoids re-dumping the full segment desc on - // every peer fetch — the previous behavior multiplied dump cost by the - // number of concurrent peer RPCs and dominated remote getRemote latency. + // invalidated by updateLocal() on every publication. This avoids + // re-dumping the full segment desc on every peer fetch — the previous + // behavior multiplied dump cost by the number of concurrent peer RPCs and + // dominated remote getRemote latency. std::shared_ptr getLocalDumpedJson(); Status synchronizeLocal(); @@ -153,7 +221,25 @@ class SegmentManager { std::atomic version_; + // Current published snapshot of the local SegmentDesc. Replaced wholesale + // by updateLocal(); never mutated in place. local_desc_lock_ only guards + // the pointer swap, not the pointee. std::shared_mutex (rather than the + // in-tree RWSpinlock) keeps the synchronization visible to + // ThreadSanitizer. + std::shared_mutex local_desc_lock_; SegmentDescRef local_desc_; + // Serializes clone-mutate-publish cycles in updateLocal(). + std::mutex local_update_mu_; + // Serializes {snapshot, putSegmentDesc} pairs in synchronizeLocal() so a + // stale in-flight put cannot overwrite a newer snapshot in the registry. + std::mutex local_sync_mu_; + // Publication counter; invalidates the per-thread snapshot caches in + // getLocal() and tags local_json_cache_ so that a slow JSON dump of an + // old snapshot can never overwrite the cache entry of a newer one. + std::atomic local_desc_version_{0}; + // Process-unique id for the thread-local cache key in getLocal(). + const uint64_t manager_id_; + ThreadLocalStorage tl_remote_cache_; std::unique_ptr registry_; @@ -164,9 +250,12 @@ class SegmentManager { std::shared_ptr subscribers_lock_; std::shared_ptr> subscribers_; - // Cache for the serialized JSON of local_desc_. Reset by synchronizeLocal. + // Cache for the serialized JSON of local_desc_. Invalidated by + // updateLocal(); local_json_cache_version_ records which publication the + // cached string was computed from. std::mutex local_json_cache_mu_; std::shared_ptr local_json_cache_; + uint64_t local_json_cache_version_{0}; }; } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h index 467a997565..a32aec8d27 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h @@ -29,13 +29,18 @@ #include #include "tent/runtime/segment.h" +#include "tent/runtime/segment_manager.h" namespace mooncake { namespace tent { +// Maintains the buffer list of the local SegmentDesc (ref-counted +// registration / deregistration). All mutations are applied through +// SegmentManager::updateLocal(), so concurrent readers of getLocal() always +// observe consistent snapshots; SegmentTracker itself holds no state besides +// the manager reference. class SegmentTracker { public: - SegmentTracker(const SegmentDescRef& local_desc) - : local_desc_(local_desc) {} + explicit SegmentTracker(SegmentManager& manager) : manager_(manager) {} ~SegmentTracker() {} @@ -43,9 +48,6 @@ class SegmentTracker { SegmentTracker& operator==(const SegmentTracker&) = delete; public: - Status query(uint64_t base, size_t length, - std::vector& result); - Status addInBatch(std::vector& desc_list, std::function&)> callback); @@ -55,13 +57,15 @@ class SegmentTracker { Status remove(uint64_t base, size_t length, std::function callback); - Status forEach(std::function callback); + // Iterates over the current snapshot; entries are immutable. Callers + // needing a mutable copy (e.g. transports scrubbing keys during + // deregistration) must copy explicitly. + Status forEach(std::function callback); private: - SegmentDescRef local_desc_; - std::mutex mutex_; + SegmentManager& manager_; }; } // namespace tent } // namespace mooncake -#endif // SEGMENT_TRACKER_H \ No newline at end of file +#endif // SEGMENT_TRACKER_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h index b8ba52d78d..9b0f635daf 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -76,6 +76,9 @@ class Workers { private: struct RouteHint { + // Owning reference to the segment snapshot; keeps all raw pointers + // below valid for the lifetime of this hint. + SegmentDescRef pin; SegmentDesc *segment; BufferDesc *buffer; const Topology::MemEntry *topo_entry; diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index d82fbb5410..ca5facb0e6 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -249,7 +249,7 @@ void ControlService::onSendData(const std::string_view& request, return; } XferDataDesc* desc = (XferDataDesc*)request.data(); - auto local_desc = manager_->getLocal().get(); + auto local_desc = manager_->getLocal(); auto peer_mem_addr = le64toh(desc->peer_mem_addr); auto length = le64toh(desc->length); @@ -273,7 +273,7 @@ void ControlService::onRecvData(const std::string_view& request, return; } XferDataDesc* desc = (XferDataDesc*)request.data(); - auto local_desc = manager_->getLocal().get(); + auto local_desc = manager_->getLocal(); auto peer_mem_addr = le64toh(desc->peer_mem_addr); auto length = le64toh(desc->length); diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp index 98f6331afb..6181171900 100644 --- a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp @@ -25,8 +25,18 @@ namespace mooncake { namespace tent { +namespace { +uint64_t nextManagerId() { + static std::atomic counter{0}; + return counter.fetch_add(1, std::memory_order_relaxed); +} +} // namespace + SegmentManager::SegmentManager(std::unique_ptr agent) - : next_id_(1), version_(0), registry_(std::move(agent)) { + : next_id_(1), + version_(0), + manager_id_(nextManagerId()), + registry_(std::move(agent)) { local_desc_ = std::make_shared(); subscribers_lock_ = std::make_shared(); subscribers_ = std::make_shared>(); @@ -60,7 +70,39 @@ Status SegmentManager::closeRemote(SegmentID handle) { return Status::OK(); } +Status SegmentManager::updateLocal( + const std::function &mutator) { + std::lock_guard g(local_update_mu_); + // No concurrent writer exists (serialized by local_update_mu_), so + // reading local_desc_ without local_desc_lock_ is safe here. + auto next = std::make_shared(*local_desc_); + CHECK_STATUS(mutator(*next)); + { + std::unique_lock guard(local_desc_lock_); + local_desc_ = std::move(next); + } + // Release pairs with the acquire load in getLocal(): a reader that + // observes the new version must also observe the pointer swap above. + // With a relaxed bump, on weakly-ordered architectures the version + // store may become visible before the swap, letting a reader tag the + // OLD snapshot with the NEW version — which its thread-local cache + // would then serve until the next publication. + local_desc_version_.fetch_add(1, std::memory_order_release); + { + std::lock_guard jg(local_json_cache_mu_); + local_json_cache_.reset(); + } + return Status::OK(); +} + Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { + SegmentDescRef ref; + CHECK_STATUS(getRemoteCached(ref, handle)); + desc = ref.get(); + return Status::OK(); +} + +Status SegmentManager::getRemoteCached(SegmentDescRef &desc, SegmentID handle) { auto &cache = tl_remote_cache_.get(); auto current_ts = getCurrentTimeInNano(); auto current_version = version_.load(std::memory_order_relaxed); @@ -78,7 +120,7 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { cache.id_to_desc_map[handle] = desc_ref; std::string peer_rpc_addr = desc_ref->rpc_server_addr; - std::string local_rpc_addr = local_desc_->rpc_server_addr; + std::string local_rpc_addr = getLocal()->rpc_server_addr; if (!peer_rpc_addr.empty() && !local_rpc_addr.empty()) { // Send a subscription request to enable proactive cache // invalidation. This is a best-effort mechanism to reduce stale @@ -92,7 +134,7 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { << "'."; } } - desc = cache.id_to_desc_map[handle].get(); + desc = cache.id_to_desc_map[handle]; assert(desc); return Status::OK(); } @@ -164,7 +206,7 @@ Status SegmentManager::makeFileRemote(SegmentDescRef &desc, desc = std::make_shared(); desc->name = segment_name; desc->type = SegmentType::File; - desc->machine_id = local_desc_->machine_id; + desc->machine_id = getLocal()->machine_id; FileSegmentDesc detail; FileBufferDesc buffer; buffer.path = path; @@ -176,26 +218,42 @@ Status SegmentManager::makeFileRemote(SegmentDescRef &desc, } std::shared_ptr SegmentManager::getLocalDumpedJson() { + // Capture the snapshot together with its publication version so a slow + // dump of an older snapshot can never overwrite the cache entry computed + // from a newer one. Acquire keeps the tag conservative: the snapshot + // read below is then guaranteed to be at least as new as the tag. + auto version = local_desc_version_.load(std::memory_order_acquire); + auto snapshot = getLocal(); { std::lock_guard g(local_json_cache_mu_); - if (local_json_cache_) return local_json_cache_; + if (local_json_cache_ && local_json_cache_version_ == version) + return local_json_cache_; } - json j = *local_desc_; + json j = *snapshot; auto computed = std::make_shared(j.dump()); std::lock_guard g(local_json_cache_mu_); - if (!local_json_cache_) { + // Store only if no publication happened since we sampled `version`; + // otherwise this dump is already outdated and must not evict a cache + // entry computed from a newer snapshot. + if (local_desc_version_.load(std::memory_order_relaxed) == version && + !local_json_cache_) { local_json_cache_ = computed; + local_json_cache_version_ = version; } - return local_json_cache_; + return computed; } Status SegmentManager::synchronizeLocal() { { - std::lock_guard g(local_json_cache_mu_); - local_json_cache_.reset(); + // Serialize {snapshot, put} pairs: without this, a put carrying an + // older snapshot could complete after (and overwrite) one carrying a + // newer snapshot in the registry, hiding a completed registration + // from peers until the next synchronizeLocal call. + std::lock_guard g(local_sync_mu_); + auto snapshot = getLocal(); + CHECK_STATUS(registry_->putSegmentDesc(snapshot)); } - CHECK_STATUS(registry_->putSegmentDesc(local_desc_)); std::vector subscribers_snapshot; { @@ -214,7 +272,7 @@ Status SegmentManager::synchronizeLocal() { // Remove subscribers that have failed (e.g., peer might shutdown) // to avoid repeated RPC failures. ControlClient::notifySegmentUpdatedAsync( - subscriber, local_desc_->name, + subscriber, getLocal()->name, /* on_failure */ [subscribers = subscribers_, lock = subscribers_lock_, subscriber] { RWSpinlock::WriteGuard guard(*lock); @@ -225,7 +283,7 @@ Status SegmentManager::synchronizeLocal() { } Status SegmentManager::deleteLocal() { - return registry_->deleteSegmentDesc(local_desc_->name); + return registry_->deleteSegmentDesc(getLocal()->name); } } // namespace tent diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp index 2f84486d91..0faf27ba25 100644 --- a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp @@ -25,39 +25,47 @@ namespace mooncake { namespace tent { -Status SegmentTracker::query(uint64_t base, size_t length, - std::vector& result) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - assert(length); +namespace { +void sortBuffers(std::vector& buffers) { + std::sort(buffers.begin(), buffers.end(), + [](const BufferDesc& lhs, const BufferDesc& rhs) -> bool { + if (lhs.addr < rhs.addr) return true; + if (lhs.addr > rhs.addr) return false; + return lhs.length > rhs.length; // prefer large interval + }); +} + +bool containsBuffer(const SegmentDesc& desc, uint64_t base, size_t length) { + auto& detail = std::get(desc.detail); for (auto& buf : detail.buffers) { - if (buf.addr > base) continue; - if (buf.addr + buf.length <= base) break; - result.push_back(&buf); - if (buf.addr + buf.length >= base + length) { - return Status::OK(); - } else { - auto new_base = buf.addr + buf.length; - auto new_length = base + length - new_base; - return query(new_base, new_length, result); - } + if (buf.addr == base && buf.length == length) return true; } - return Status::InvalidArgument("Some buffers are not registered"); + return false; } +} // namespace Status SegmentTracker::add(uint64_t base, size_t length, std::function callback) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - mutex_.lock(); - for (auto& buf : detail.buffers) { - if (buf.addr == base && buf.length == length) { - buf.ref_count++; - mutex_.unlock(); + // Read-only pre-scan on the current snapshot: the common miss path pays + // no clone/publication. A hit is re-verified under the writer mutex; a + // racing insert of the same range degenerates to today's benign + // duplicate-registration behavior. + if (containsBuffer(*manager_.getLocal(), base, length)) { + bool found = false; + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto& buf : detail.buffers) { + if (buf.addr == base && buf.length == length) { + buf.ref_count++; + found = true; + break; + } + } return Status::OK(); - } + })); + if (found) return Status::OK(); } - mutex_.unlock(); BufferDesc new_desc; new_desc.addr = base; new_desc.length = length; @@ -71,86 +79,129 @@ Status SegmentTracker::add(uint64_t base, size_t length, new_desc.ref_count = 1; auto status = callback(new_desc); if (!status.ok()) return status; - mutex_.lock(); - detail.buffers.push_back(new_desc); - std::sort(detail.buffers.begin(), detail.buffers.end(), - [](const BufferDesc& lhs, BufferDesc& rhs) -> bool { - if (lhs.addr < rhs.addr) return true; - if (lhs.addr > rhs.addr) return false; - return lhs.length > rhs.length; // prefer large interval - }); - mutex_.unlock(); - return Status::OK(); + return manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + detail.buffers.push_back(std::move(new_desc)); + sortBuffers(detail.buffers); + return Status::OK(); + }); } Status SegmentTracker::addInBatch( std::vector& desc_list, std::function&)> callback) { std::vector new_desc_list; - for (auto& desc : desc_list) { - bool found = false; - { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - for (auto& buf : detail.buffers) { - if (buf.addr == desc.addr && buf.length == desc.length) { - buf.ref_count++; - found = true; - break; - } + // Read-only pre-scan (see add()): skip the ref-count publication when no + // entry duplicates an already-registered range. + bool any_dup = false; + { + auto snapshot = manager_.getLocal(); + for (auto& entry : desc_list) { + if (containsBuffer(*snapshot, entry.addr, entry.length)) { + any_dup = true; + break; } } - if (!found) new_desc_list.push_back(std::move(desc)); + } + // Ranges whose ref_count we bumped; used to roll back if the callback + // fails. The bump must happen under the writer mutex *before* the + // callback: it is what pins the duplicate entry (ref_count >= 2) so a + // concurrent unregister cannot erase it out from under this + // registration. + std::vector> bumped; + if (any_dup) { + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto& entry : desc_list) { + bool found = false; + for (auto& buf : detail.buffers) { + if (buf.addr == entry.addr && buf.length == entry.length) { + buf.ref_count++; + found = true; + break; + } + } + if (found) { + bumped.emplace_back(entry.addr, entry.length); + } else { + new_desc_list.push_back(std::move(entry)); + } + } + return Status::OK(); + })); + } else { + new_desc_list = std::move(desc_list); } auto status = callback(new_desc_list); - if (!status.ok()) return status; - { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); + if (!status.ok()) { + // Roll back the duplicate ref-counts so a failed registration does + // not leave buffers pinned forever. + if (!bumped.empty()) { + manager_.updateLocal([&](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + for (auto& range : bumped) { + for (auto it = detail.buffers.begin(); + it != detail.buffers.end(); ++it) { + if (it->addr == range.first && + it->length == range.second) { + it->ref_count--; + // The original owner unregistered while we held + // the extra reference; drop the entry so it is + // no longer advertised. + if (it->ref_count == 0) detail.buffers.erase(it); + break; + } + } + } + return Status::OK(); + }); + } + return status; + } + return manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); for (auto& new_desc : new_desc_list) { detail.buffers.push_back(new_desc); } - std::sort(detail.buffers.begin(), detail.buffers.end(), - [](const BufferDesc& lhs, BufferDesc& rhs) -> bool { - if (lhs.addr < rhs.addr) return true; - if (lhs.addr > rhs.addr) return false; - return lhs.length > rhs.length; // prefer large interval - }); - } - return Status::OK(); + sortBuffers(detail.buffers); + return Status::OK(); + }); } Status SegmentTracker::remove(uint64_t base, size_t length, std::function callback) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - mutex_.lock(); - for (auto it = detail.buffers.begin(); it != detail.buffers.end(); ++it) { - if (it->addr == base && (!length || it->length == length)) { - it->ref_count--; - Status status = Status::OK(); - if (it->ref_count == 0) { - BufferDesc clone = *it; - detail.buffers.erase(it); - mutex_.unlock(); - status = callback(clone); - } else { - mutex_.unlock(); + bool removed = false; + BufferDesc removed_desc; + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto it = detail.buffers.begin(); it != detail.buffers.end(); + ++it) { + if (it->addr == base && (!length || it->length == length)) { + it->ref_count--; + if (it->ref_count == 0) { + removed_desc = *it; + detail.buffers.erase(it); + removed = true; + } + break; } - return status; } - } - mutex_.unlock(); + return Status::OK(); + })); + if (removed) return callback(removed_desc); return Status::OK(); } -Status SegmentTracker::forEach(std::function callback) { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - for (auto& buf : detail.buffers) { +Status SegmentTracker::forEach( + std::function callback) { + auto snapshot = manager_.getLocal(); + assert(snapshot->type == SegmentType::Memory); + for (const auto& buf : + std::get(snapshot->detail).buffers) { auto status = callback(buf); if (!status.ok()) return status; } diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index b1f2116eee..12bfbb8abb 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -265,14 +265,16 @@ std::string getMachineID() { Status TransferEngineImpl::setupLocalSegment() { auto& manager = metadata_->segmentManager(); - auto segment = manager.getLocal(); - segment->name = local_segment_name_; - segment->type = SegmentType::Memory; - segment->machine_id = getMachineID(); - segment->rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); - auto& detail = std::get(segment->detail); - detail.topology = *(topology_.get()); - local_segment_tracker_ = std::make_unique(segment); + CHECK_STATUS(manager.updateLocal([&](SegmentDesc& segment) -> Status { + segment.name = local_segment_name_; + segment.type = SegmentType::Memory; + segment.machine_id = getMachineID(); + segment.rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); + auto& detail = std::get(segment.detail); + detail.topology = *(topology_.get()); + return Status::OK(); + })); + local_segment_tracker_ = std::make_unique(manager); return manager.synchronizeLocal(); } @@ -430,10 +432,13 @@ Status TransferEngineImpl::deconstruct() { staging_proxy_.reset(); if (local_segment_tracker_) { - local_segment_tracker_->forEach([&](BufferDesc& desc) -> Status { + local_segment_tracker_->forEach([&](const BufferDesc& desc) -> Status { + // Snapshot entries are immutable; transports may scrub fields of + // their deregistration argument, so hand them a copy. + BufferDesc copy = desc; for (size_t type = 0; type < kSupportedTransportTypes; ++type) { if (transport_list_[type]) - transport_list_[type]->removeMemoryBuffer(desc); + transport_list_[type]->removeMemoryBuffer(copy); } return Status::OK(); }); @@ -514,9 +519,10 @@ Status TransferEngineImpl::closeSegment(SegmentID handle) { } Status TransferEngineImpl::getSegmentInfo(SegmentID handle, SegmentInfo& info) { - SegmentDesc* desc = nullptr; + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef desc; if (handle == LOCAL_SEGMENT_ID) { - desc = metadata_->segmentManager().getLocal().get(); + desc = metadata_->segmentManager().getLocal(); } else { CHECK_STATUS(metadata_->segmentManager().getRemoteCached(desc, handle)); } @@ -936,9 +942,10 @@ Status TransferEngineImpl::validateTransportHint(const Request& req, SelectionResult TransferEngineImpl::getTransportType(const Request& request, int transport_index) { - SegmentDesc* desc; + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef desc; if (request.target_id == LOCAL_SEGMENT_ID) { - desc = metadata_->segmentManager().getLocal().get(); + desc = metadata_->segmentManager().getLocal(); } else { auto status = metadata_->segmentManager().getRemoteCached( desc, request.target_id); @@ -1225,7 +1232,8 @@ std::vector resolveRequestBoundaries( // Group requests by target_id so withCachedSegment fires at most once per // peer. std::vector boundaries(requests.size()); - auto* local_desc = metadata->segmentManager().getLocal().get(); + // Owning reference: keeps the snapshot alive while we read through it. + auto local_desc = metadata->segmentManager().getLocal(); if (local_desc) { for (size_t i = 0; i < requests.size(); ++i) { @@ -1285,8 +1293,10 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, SegmentDesc* desc = nullptr; BufferDesc* entry = nullptr; + // Owning reference: `entry` is used after the lambda returns. + SegmentDescRef pin; auto status = metadata_->segmentManager().withCachedSegment( - request.target_id, [&](SegmentDesc* segment) { + request.target_id, pin, [&](SegmentDesc* segment) { desc = segment; entry = desc->findBuffer(request.target_offset, request.length); if (!entry) diff --git a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp index d50fed7fef..99f7be72c0 100644 --- a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp @@ -160,9 +160,12 @@ Status AscendDirectTransport::initHixl(const std::shared_ptr &conf) { auto hixl_name = host_ip + ":" + std::to_string(port); local_hixl_name_ = hixl_name; - auto segment = metadata_->segmentManager().getLocal(); - auto &detail = std::get(segment->detail); - detail.device_attrs["hixl_name"] = hixl_name; + CHECK_STATUS(metadata_->segmentManager().updateLocal( + [&](SegmentDesc &segment) -> Status { + auto &detail = std::get(segment.detail); + detail.device_attrs["hixl_name"] = hixl_name; + return Status::OK(); + })); hixl_ = std::make_unique(); if (!hixl_) return Status::InternalError("Create hixl failed."); diff --git a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp index 046b253be4..ca5b8702b6 100644 --- a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp @@ -174,7 +174,8 @@ Status MnnvlTransport::submitTransferTasks( // Get local segment for buffer lookup auto &segment_manager = metadata_->segmentManager(); - SegmentDesc *local_segment = segment_manager.getLocal().get(); + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef local_segment = segment_manager.getLocal(); if (!local_segment) return Status::InternalError("Local segment not found" LOC_MARK); @@ -553,9 +554,11 @@ Status MnnvlTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, RWSpinlock::WriteGuard guard(relocate_lock_); BufferDesc *buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto &segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc *segment) { + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc *segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->mnnvl_handle.empty()) return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp index 70707beea8..157c97f073 100644 --- a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp @@ -134,7 +134,8 @@ Status NVLinkTransport::submitTransferTasks( // Get local segment for buffer lookup auto& segment_manager = metadata_->segmentManager(); - SegmentDesc* local_segment = segment_manager.getLocal().get(); + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef local_segment = segment_manager.getLocal(); if (!local_segment) return Status::InternalError("Local segment not found" LOC_MARK); @@ -459,9 +460,11 @@ Status NVLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr, RWSpinlock::WriteGuard guard(relocate_lock_); BufferDesc* buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto& segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc* segment) { + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc* segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->shm_path.empty()) return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index 741181441a..93736e24dd 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -543,22 +543,23 @@ Status RdmaTransport::removeMemoryBuffer(BufferDesc& desc) { Status RdmaTransport::setupLocalSegment() { auto& manager = metadata_->segmentManager(); - auto segment = manager.getLocal(); - assert(segment); - // Store RDMA server name for dual-NIC setups; when it differs from - // local_segment_name_ the peer will use it for NIC path construction. - if (rdma_server_name_ != local_segment_name_) { - segment->rdma_server_name = rdma_server_name_; - } - auto& detail = std::get(segment->detail); - for (auto& context : context_set_) { - if (context->status() != RdmaContext::DEVICE_ENABLED) continue; - DeviceDesc device_desc; - device_desc.name = context->name(); - device_desc.lid = context->lid(); - device_desc.gid = context->gid(); - detail.devices.push_back(device_desc); - } + CHECK_STATUS(manager.updateLocal([&](SegmentDesc& segment) -> Status { + // Store RDMA server name for dual-NIC setups; when it differs from + // local_segment_name_ the peer will use it for NIC path construction. + if (rdma_server_name_ != local_segment_name_) { + segment.rdma_server_name = rdma_server_name_; + } + auto& detail = std::get(segment.detail); + for (auto& context : context_set_) { + if (context->status() != RdmaContext::DEVICE_ENABLED) continue; + DeviceDesc device_desc; + device_desc.name = context->name(); + device_desc.lid = context->lid(); + device_desc.gid = context->gid(); + detail.devices.push_back(device_desc); + } + return Status::OK(); + })); return manager.synchronizeLocal(); } @@ -606,13 +607,11 @@ int RdmaTransport::onSetupRdmaConnections(const BootstrapDesc& peer_desc, std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, int device_id) { - SegmentDesc* segment_desc = nullptr; - std::string rpc_server_addr, target_seg_name, target_dev_name; + std::string rpc_server_addr, target_seg_name, target_dev_name, + target_nic_path_name; auto status = metadata_->segmentManager().withCachedSegment( target_id, [&](SegmentDesc* segment) { - segment_desc = segment; - if (segment->type != SegmentType::Memory) { return Status::NeedsRefreshCache( "Segment type is not Memory" LOC_MARK); @@ -624,6 +623,7 @@ std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, auto topo = &std::get(segment->detail).topology; target_seg_name = segment->name; + target_nic_path_name = segment->nicPathServerName(); target_dev_name = topo->getNicName(device_id); if (target_seg_name.empty() || target_dev_name.empty()) { return Status::NeedsRefreshCache( @@ -642,8 +642,7 @@ std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, return nullptr; } std::shared_ptr endpoint; - std::string peer_name = - MakeNicPath(segment_desc->nicPathServerName(), target_dev_name); + std::string peer_name = MakeNicPath(target_nic_path_name, target_dev_name); endpoint = context->endpointStore()->getOrInsert(peer_name); if (!endpoint) { LOG(ERROR) << "Cannot allocate endpoint " << peer_name; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 500cc9e3e4..319fec32fc 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -260,8 +260,8 @@ std::shared_ptr Workers::getEndpoint(Workers::PostPath path) { auto target_id = path.remote_segment_id; auto device_id = path.remote_device_id; - auto status = - segment_manager.withCachedSegment(target_id, [&](SegmentDesc* segment) { + auto status = segment_manager.withCachedSegment( + target_id, hint.pin, [&](SegmentDesc* segment) { hint.segment = segment; if (segment->type != SegmentType::Memory) { return Status::NeedsRefreshCache( @@ -670,7 +670,7 @@ Status Workers::getRouteHint(RouteHint& hint, SegmentID segment_id, uint64_t addr, uint64_t length) { auto& segment_manager = transport_->metadata_->segmentManager(); CHECK_STATUS(segment_manager.withCachedSegment( - segment_id, [&](SegmentDesc* segment) { + segment_id, hint.pin, [&](SegmentDesc* segment) { hint.segment = segment; hint.buffer = segment->findBuffer(addr, length); if (!hint.buffer) diff --git a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp index ccd737eeb6..55f4a0d20f 100644 --- a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp @@ -263,9 +263,11 @@ Status ShmTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, RWSpinlock::WriteGuard guard(relocate_lock_); BufferDesc *buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto &segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc *segment) { + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc *segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->shm_path.empty()) return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index fa70ae824d..7329976706 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -40,6 +40,16 @@ target_include_directories(tent_coalesce_regions_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_coalesce_regions_test COMMAND tent_coalesce_regions_test) +add_executable(segment_manager_test segment_manager_test.cpp) +target_link_libraries(segment_manager_test PRIVATE gtest gtest_main + tent_link_group) +if(TARGET asio_shared) + target_link_libraries(segment_manager_test PRIVATE asio_shared) +endif() +target_include_directories(segment_manager_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME segment_manager_test COMMAND segment_manager_test) + add_executable(request_merge_test request_merge_test.cpp) target_link_libraries(request_merge_test PRIVATE gtest gtest_main tent_link_group) diff --git a/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp b/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp new file mode 100644 index 0000000000..77bb31d7fc --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp @@ -0,0 +1,309 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression tests for issue #2477: concurrent registerLocalMemory / +// unregisterLocalMemory racing with lock-free readers of the local +// SegmentDesc. The local desc is published as immutable copy-on-write +// snapshots; these tests assert the snapshot semantics (readers never +// observe torn or unsorted buffer lists) and, when built with +// -fsanitize=thread, additionally prove the absence of data races between +// SegmentTracker writers and getLocal() / getLocalDumpedJson() readers. +// A port of ConcurrentWritersVsSnapshotReaders to the pre-fix in-place +// mutation API fails its invariants on stock builds and reports data races +// under TSAN. + +#include + +#include +#include +#include +#include +#include + +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/segment_registry.h" +#include "tent/runtime/segment_tracker.h" + +namespace mooncake { +namespace tent { + +namespace { + +std::unique_ptr makeManager() { + // No registry: these tests never touch remote segments or + // synchronizeLocal(). + auto manager = std::make_unique(nullptr); + EXPECT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + desc.name = "local_test_segment"; + desc.type = SegmentType::Memory; + desc.machine_id = "test_machine"; + return Status::OK(); + }) + .ok()); + return manager; +} + +// Location derived from the buffer address; readers use it to detect +// partially-constructed or torn BufferDesc entries. +std::string locationFor(uint64_t addr) { + return "cpu:" + std::to_string(addr % 4096); +} + +BufferDesc makeBuffer(uint64_t addr, uint64_t length) { + BufferDesc desc; + desc.addr = addr; + desc.length = length; + desc.location = locationFor(addr); + desc.ref_count = 1; + return desc; +} + +const std::vector& buffersOf(const SegmentDescRef& snapshot) { + return std::get(snapshot->detail).buffers; +} + +} // namespace + +TEST(SegmentManagerTest, UpdateLocalPublishesImmutableSnapshots) { + auto manager = makeManager(); + auto before = manager->getLocal(); + ASSERT_EQ(before->name, "local_test_segment"); + ASSERT_TRUE(buffersOf(before).empty()); + + ASSERT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + detail.buffers.push_back(makeBuffer(0x1000, 0x1000)); + return Status::OK(); + }) + .ok()); + + // The old snapshot is untouched; the new snapshot sees the mutation. + EXPECT_TRUE(buffersOf(before).empty()); + auto after = manager->getLocal(); + ASSERT_EQ(buffersOf(after).size(), 1u); + EXPECT_EQ(buffersOf(after)[0].addr, 0x1000u); + EXPECT_NE(before.get(), after.get()); +} + +TEST(SegmentManagerTest, UpdateLocalFailureDoesNotPublish) { + auto manager = makeManager(); + auto before = manager->getLocal(); + auto status = manager->updateLocal([](SegmentDesc& desc) -> Status { + desc.name = "must_not_be_published"; + return Status::InvalidArgument("injected failure"); + }); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(manager->getLocal().get(), before.get()); + EXPECT_EQ(manager->getLocal()->name, "local_test_segment"); +} + +TEST(SegmentManagerTest, JsonCacheInvalidatedOnPublication) { + auto manager = makeManager(); + auto dump1 = manager->getLocalDumpedJson(); + auto dump2 = manager->getLocalDumpedJson(); + EXPECT_EQ(dump1.get(), dump2.get()); // served from cache + + ASSERT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + detail.buffers.push_back(makeBuffer(0x2000, 0x1000)); + return Status::OK(); + }) + .ok()); + + auto dump3 = manager->getLocalDumpedJson(); + EXPECT_NE(*dump1, *dump3); + EXPECT_NE(dump3->find("8192"), std::string::npos); // 0x2000 serialized +} + +TEST(SegmentTrackerTest, RefCountedAddRemove) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + auto noop = [](std::vector&) -> Status { return Status::OK(); }; + std::vector first{makeBuffer(0x10000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(first, noop).ok()); + std::vector second{makeBuffer(0x10000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(second, noop).ok()); + + auto snapshot = manager->getLocal(); + ASSERT_EQ(buffersOf(snapshot).size(), 1u); + EXPECT_EQ(buffersOf(snapshot)[0].ref_count, 2); + + int remove_callbacks = 0; + auto on_remove = [&](BufferDesc&) -> Status { + remove_callbacks++; + return Status::OK(); + }; + ASSERT_TRUE(tracker.remove(0x10000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 0); // still referenced + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + + ASSERT_TRUE(tracker.remove(0x10000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 1); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +TEST(SegmentTrackerTest, AddInBatchCallbackFailureRollsBackRefCounts) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + auto ok = [](std::vector&) -> Status { return Status::OK(); }; + std::vector first{makeBuffer(0x20000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(first, ok).ok()); + ASSERT_EQ(buffersOf(manager->getLocal())[0].ref_count, 1); + + // A duplicate registration whose transport callback fails must not leave + // the ref-count bumped, or the buffer could never be deregistered. + auto fail = [](std::vector&) -> Status { + return Status::InternalError("injected transport failure"); + }; + std::vector dup{makeBuffer(0x20000, 0x1000)}; + EXPECT_FALSE(tracker.addInBatch(dup, fail).ok()); + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + EXPECT_EQ(buffersOf(manager->getLocal())[0].ref_count, 1); + + int remove_callbacks = 0; + auto on_remove = [&](BufferDesc&) -> Status { + remove_callbacks++; + return Status::OK(); + }; + ASSERT_TRUE(tracker.remove(0x20000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 1); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +TEST(SegmentTrackerTest, AddProbesRealMemoryAndRefCounts) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + constexpr size_t kSize = 1 << 20; + void* mem = malloc(kSize); + ASSERT_NE(mem, nullptr); + auto base = reinterpret_cast(mem); + + int add_callbacks = 0; + auto on_add = [&](BufferDesc& desc) -> Status { + add_callbacks++; + EXPECT_EQ(desc.addr, base); + EXPECT_FALSE(desc.location.empty()); // NUMA probe ran + return Status::OK(); + }; + ASSERT_TRUE(tracker.add(base, kSize, on_add).ok()); + EXPECT_EQ(add_callbacks, 1); + // Re-adding the same range takes the ref-count fast path: no new probe. + ASSERT_TRUE(tracker.add(base, kSize, on_add).ok()); + EXPECT_EQ(add_callbacks, 1); + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + EXPECT_EQ(buffersOf(manager->getLocal())[0].ref_count, 2); + + auto noop = [](BufferDesc&) -> Status { return Status::OK(); }; + ASSERT_TRUE(tracker.remove(base, kSize, noop).ok()); + ASSERT_TRUE(tracker.remove(base, kSize, noop).ok()); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); + free(mem); +} + +// The actual #2477 regression: register/unregister churn concurrent with +// lock-free snapshot readers. With in-place mutation this is a data race on +// MemorySegmentDesc::buffers (vector push_back/erase/sort vs. iteration): +// ThreadSanitizer reports it and the location invariant below can observe +// partially-constructed entries. With copy-on-write snapshots every reader +// observes a fully-consistent (possibly stale) buffer list. +TEST(SegmentTrackerTest, ConcurrentWritersVsSnapshotReaders) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + constexpr int kWriters = 4; + constexpr int kReaders = 4; + constexpr int kIterations = 300; + constexpr int kBuffersPerBatch = 8; + constexpr uint64_t kLength = 0x1000; + + std::atomic done{false}; + std::atomic failures{0}; + + auto noop = [](std::vector&) -> Status { return Status::OK(); }; + + std::vector writers; + writers.reserve(kWriters); + for (int w = 0; w < kWriters; ++w) { + writers.emplace_back([&, w] { + // Writers 0 and 1 share an address range so ref-count bumps, + // duplicate-registration races and erase-vs-bump interleavings + // are exercised concurrently, not just disjoint inserts. + const uint64_t base = (w < 2 ? 1 : w + 1) * 0x100000000ULL; + for (int iter = 0; iter < kIterations; ++iter) { + std::vector batch; + batch.reserve(kBuffersPerBatch); + for (int i = 0; i < kBuffersPerBatch; ++i) { + batch.push_back( + makeBuffer(base + i * kLength * 2, kLength)); + } + if (!tracker.addInBatch(batch, noop).ok()) failures++; + for (int i = 0; i < kBuffersPerBatch; ++i) { + auto on_remove = [](BufferDesc&) -> Status { + return Status::OK(); + }; + if (!tracker + .remove(base + i * kLength * 2, kLength, on_remove) + .ok()) + failures++; + } + } + }); + } + + std::vector readers; + readers.reserve(kReaders); + for (int r = 0; r < kReaders; ++r) { + readers.emplace_back([&, r] { + uint64_t rounds = 0; + while (!done.load(std::memory_order_acquire)) { + auto snapshot = manager->getLocal(); + const auto& buffers = buffersOf(snapshot); + uint64_t prev_addr = 0; + for (const auto& buf : buffers) { + // Entries must be fully constructed and sorted; a torn + // read of an in-place mutated vector violates these. + if (buf.length != kLength || + buf.location != locationFor(buf.addr) || + buf.addr < prev_addr) { + failures++; + } + prev_addr = buf.addr; + } + // Exercise findBuffer() through the snapshot as transports + // do, and the JSON dump path peers hit via GetSegmentDesc. + snapshot->findBuffer(0x100000000ULL, kLength); + if (rounds++ % 64 == 0) { + auto dump = manager->getLocalDumpedJson(); + if (!dump || dump->empty()) failures++; + } + } + }); + } + + for (auto& t : writers) t.join(); + done.store(true, std::memory_order_release); + for (auto& t : readers) t.join(); + + EXPECT_EQ(failures.load(), 0); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +} // namespace tent +} // namespace mooncake From bd8cc9f95ea80c2df82df47e8475f7c39941346d Mon Sep 17 00:00:00 2001 From: Aoi Date: Sun, 5 Jul 2026 22:53:40 +0800 Subject: [PATCH 09/57] [Store] rust: format mooncake store bindings (#2737) Co-authored-by: Aionw --- mooncake-store/go/mooncakestore/errors.go | 22 ++++----- mooncake-store/rust/build.rs | 24 ++++++---- mooncake-store/rust/examples/basic_usage.rs | 15 ++++--- mooncake-store/rust/src/store.rs | 50 +++++++++------------ 4 files changed, 56 insertions(+), 55 deletions(-) diff --git a/mooncake-store/go/mooncakestore/errors.go b/mooncake-store/go/mooncakestore/errors.go index 6c673e281f..0fee47eab6 100644 --- a/mooncake-store/go/mooncakestore/errors.go +++ b/mooncake-store/go/mooncakestore/errors.go @@ -17,18 +17,18 @@ package mooncakestore import "errors" var ( - ErrStoreNil = errors.New("mooncakestore: store handle is nil") - ErrSetupFailed = errors.New("mooncakestore: setup failed") - ErrInitAllFailed = errors.New("mooncakestore: init_all failed") - ErrHealthCheck = errors.New("mooncakestore: health check failed") - ErrPut = errors.New("mooncakestore: put failed") - ErrGet = errors.New("mooncakestore: get failed") - ErrRemove = errors.New("mooncakestore: remove failed") - ErrExist = errors.New("mooncakestore: existence check failed") - ErrGetSize = errors.New("mooncakestore: get size failed") + ErrStoreNil = errors.New("mooncakestore: store handle is nil") + ErrSetupFailed = errors.New("mooncakestore: setup failed") + ErrInitAllFailed = errors.New("mooncakestore: init_all failed") + ErrHealthCheck = errors.New("mooncakestore: health check failed") + ErrPut = errors.New("mooncakestore: put failed") + ErrGet = errors.New("mooncakestore: get failed") + ErrRemove = errors.New("mooncakestore: remove failed") + ErrExist = errors.New("mooncakestore: existence check failed") + ErrGetSize = errors.New("mooncakestore: get size failed") ErrRegisterBuffer = errors.New("mooncakestore: register buffer failed") ErrUnregisterBuffer = errors.New("mooncakestore: unregister buffer failed") ErrBatchOp = errors.New("mooncakestore: batch operation failed") - ErrHostname = errors.New("mooncakestore: get hostname failed") - ErrInvalidArgument = errors.New("mooncakestore: invalid argument") + ErrHostname = errors.New("mooncakestore: get hostname failed") + ErrInvalidArgument = errors.New("mooncakestore: invalid argument") ) diff --git a/mooncake-store/rust/build.rs b/mooncake-store/rust/build.rs index abcd95b06c..edf3710ba0 100644 --- a/mooncake-store/rust/build.rs +++ b/mooncake-store/rust/build.rs @@ -203,19 +203,26 @@ fn main() { // common/base library (contains mooncake::Status etc.) println!( "cargo:rustc-link-search=native={}", - build_dir.join("mooncake-transfer-engine/src/common/base").display() + build_dir + .join("mooncake-transfer-engine/src/common/base") + .display() ); // CUDA runtime libraries (needed by transfer_engine RDMA transport). let cuda_home = env::var("CUDA_HOME") .or_else(|_| env::var("CUDA_PATH")) .unwrap_or_else(|_| "/usr/local/cuda".to_string()); - println!("cargo:rustc-link-search=native={}/targets/x86_64-linux/lib", cuda_home); + println!( + "cargo:rustc-link-search=native={}/targets/x86_64-linux/lib", + cuda_home + ); // cachelib_memory_allocator is a static library built alongside mooncake_store. println!( "cargo:rustc-link-search=native={}", - build_dir.join("mooncake-store/src/cachelib_memory_allocator").display() + build_dir + .join("mooncake-store/src/cachelib_memory_allocator") + .display() ); println!("cargo:rustc-link-lib=mooncake_store"); @@ -231,8 +238,8 @@ fn main() { println!("cargo:rustc-link-lib=stdc++"); println!("cargo:rustc-link-lib=glog"); println!("cargo:rustc-link-lib=gflags"); - println!("cargo:rustc-link-lib=numa"); // NUMA binding - println!("cargo:rustc-link-lib=curl"); // HTTP metadata plugin + println!("cargo:rustc-link-lib=numa"); // NUMA binding + println!("cargo:rustc-link-lib=curl"); // HTTP metadata plugin println!("cargo:rustc-link-lib=ibverbs"); // RDMA transport println!("cargo:rustc-link-lib=yaml-cpp"); // tenant quota policy connector println!("cargo:rustc-link-lib=pthread"); @@ -241,7 +248,8 @@ fn main() { // ----------------------------------------------------------------------- // Header path for bindgen // ----------------------------------------------------------------------- - let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR")); + let manifest_dir = + PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR")); let mut search_dirs = Vec::new(); let explicit_lib_dir = env::var("MOONCAKE_STORE_LIB_DIR").ok().map(PathBuf::from); @@ -365,8 +373,8 @@ fn main() { println!("cargo:rustc-link-lib=gcov"); } - let include_dir = env::var("MOONCAKE_STORE_INCLUDE_DIR") - .unwrap_or_else(|_| "../include".to_string()); + let include_dir = + env::var("MOONCAKE_STORE_INCLUDE_DIR").unwrap_or_else(|_| "../include".to_string()); let header = format!("{include_dir}/store_c.h"); diff --git a/mooncake-store/rust/examples/basic_usage.rs b/mooncake-store/rust/examples/basic_usage.rs index 4159b36978..c6f2e1b442 100644 --- a/mooncake-store/rust/examples/basic_usage.rs +++ b/mooncake-store/rust/examples/basic_usage.rs @@ -67,8 +67,8 @@ fn main() { // In CI or environments without a running metadata server the setup // call is expected to fail – this is not an error in the bindings // themselves. - let metadata_server = - std::env::var("MC_METADATA_SERVER").unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string()); + let metadata_server = std::env::var("MC_METADATA_SERVER") + .unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string()); println!("Connecting to metadata server: {metadata_server}"); @@ -78,7 +78,7 @@ fn main() { 512 << 20, // global_segment_size = 512 MiB 128 << 20, // local_buffer_size = 128 MiB "tcp", - "", // device_name (auto-select) + "", // device_name (auto-select) "127.0.0.1:50051", ) { eprintln!( @@ -98,13 +98,16 @@ fn main() { eprintln!("[FAIL] put() failed: {e}"); std::process::exit(1); } - println!("[OK] put(\"{key}\", {:?})", std::str::from_utf8(value).unwrap()); + println!( + "[OK] put(\"{key}\", {:?})", + std::str::from_utf8(value).unwrap() + ); // Step 4: Check existence. match store.is_exist(key) { - Ok(true) => println!("[OK] is_exist(\"{key}\") = true"), + Ok(true) => println!("[OK] is_exist(\"{key}\") = true"), Ok(false) => println!("[WARN] is_exist(\"{key}\") = false (unexpected)"), - Err(e) => eprintln!("[FAIL] is_exist() failed: {e}"), + Err(e) => eprintln!("[FAIL] is_exist() failed: {e}"), } // Step 5: Get size. diff --git a/mooncake-store/rust/src/store.rs b/mooncake-store/rust/src/store.rs index cb59439c7f..fdc0941d7e 100644 --- a/mooncake-store/rust/src/store.rs +++ b/mooncake-store/rust/src/store.rs @@ -57,7 +57,11 @@ impl ReplicateConfig { fn to_ffi( &self, ) -> Result< - (ffi::mooncake_replicate_config_t, Vec, Vec<*const libc::c_char>), + ( + ffi::mooncake_replicate_config_t, + Vec, + Vec<*const libc::c_char>, + ), StoreError, > { let strings: Vec = self @@ -277,8 +281,7 @@ impl MooncakeStore { pub fn get(&self, key: &str) -> Result, StoreError> { let size = self.get_size(key)?; let mut buf = vec![0u8; size as usize]; - let written = - unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? }; + let written = unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? }; buf.truncate(written as usize); Ok(buf) } @@ -345,9 +348,8 @@ impl MooncakeStore { /// read by another client. pub fn remove(&self, key: &str, force: bool) -> Result<(), StoreError> { let key_c = CString::new(key)?; - let rc = unsafe { - ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force)) - }; + let rc = + unsafe { ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force)) }; if rc != 0 { return Err(StoreError::OperationFailed(rc)); } @@ -451,8 +453,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?; let cfg_ptr = _c_config @@ -506,8 +507,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let mut results = vec![0i64; count]; @@ -528,10 +528,7 @@ impl MooncakeStore { } /// Batch check existence of multiple keys. - pub fn batch_is_exist( - &self, - keys: &[&str], - ) -> Result, StoreError> { + pub fn batch_is_exist(&self, keys: &[&str]) -> Result, StoreError> { let count = keys.len(); if count == 0 { return Ok(Vec::new()); @@ -540,8 +537,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let mut results = vec![0i32; count]; @@ -665,15 +661,13 @@ mod tests { preferred_segments: vec!["bad\0segment".to_string()], }; - assert!(matches!( - config.to_ffi(), - Err(StoreError::InvalidString(_)) - )); + assert!(matches!(config.to_ffi(), Err(StoreError::InvalidString(_)))); } #[test] fn prepare_config_none_returns_null_config() { - let (c_cfg, strings, ptrs) = MooncakeStore::prepare_config(None).expect("prepare should succeed"); + let (c_cfg, strings, ptrs) = + MooncakeStore::prepare_config(None).expect("prepare should succeed"); assert!(c_cfg.is_none()); assert!(strings.is_empty()); assert!(ptrs.is_empty()); @@ -706,7 +700,9 @@ mod tests { #[test] fn batch_is_exist_empty() { let store = MooncakeStore::new().expect("new should succeed"); - let results = store.batch_is_exist(&[]).expect("empty batch should succeed"); + let results = store + .batch_is_exist(&[]) + .expect("empty batch should succeed"); assert!(results.is_empty()); } @@ -718,10 +714,7 @@ mod tests { let sizes = &[100usize, 200]; let result = unsafe { store.batch_put_from(keys, &buffers, sizes, None) }; - assert!(matches!( - result, - Err(StoreError::InvalidArgument(_)) - )); + assert!(matches!(result, Err(StoreError::InvalidArgument(_)))); } #[test] @@ -732,9 +725,6 @@ mod tests { let sizes = &[100usize]; let result = unsafe { store.batch_get_into(keys, &buffers, sizes) }; - assert!(matches!( - result, - Err(StoreError::InvalidArgument(_)) - )); + assert!(matches!(result, Err(StoreError::InvalidArgument(_)))); } } From 4febcca55c2fec6e16cc113e1bac1993293b269e Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 6 Jul 2026 10:07:59 +0800 Subject: [PATCH 10/57] [TE] Fix MNNVL staging path reading capabilities from RDMA transport (#2751) In findStagingPolicy the "pure mnnvl" branch was guarded by transport_list_[MNNVL] && transport_list_[NVLINK] but read capabilities() from transport_list_[RDMA]. When RDMA is absent (the actual "pure mnnvl" case this branch is meant to serve), transport_list_[RDMA] is a null shared_ptr and calling capabilities() on it dereferences null. Even with RDMA present, using RDMA's caps to drive MNNVL staging decisions is semantically wrong. Read caps from transport_list_[MNNVL], matching the guard and the pattern already used by case 1 (RDMA) and case 3 (TPU). Signed-off-by: staryxchen --- .../tent/src/runtime/transfer_engine_impl.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 12bfbb8abb..a70f3da089 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -1336,7 +1336,7 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, } // case 2: pure mnnvl if (transport_list_[MNNVL] && transport_list_[NVLINK]) { - auto& xport = transport_list_[RDMA]; + auto& xport = transport_list_[MNNVL]; auto& caps = xport->capabilities(); if (local_mtype == MTYPE_CPU && remote_mtype == MTYPE_CPU && !caps.dram_to_dram) { From 712f4724249a8d52108e7d83422b3ec699f6fdcb Mon Sep 17 00:00:00 2001 From: LZW <99333079+Lin-z-w@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:20:43 +0800 Subject: [PATCH 11/57] [Store] add etcd tenant quota connector (#2687) --- .../mooncake-store-deployment-guide.md | 16 ++- docs/source/design/mooncake-store.md | 2 +- .../include/tenant_quota_policy_store.h | 19 ++- mooncake-store/src/master_service.cpp | 3 +- .../src/tenant_quota_policy_store.cpp | 119 +++++++++++++++- mooncake-store/tests/tenant_quota_test.cpp | 131 ++++++++++++++++++ 6 files changed, 277 insertions(+), 13 deletions(-) diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 5333d7d3b2..14be116665 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -278,7 +278,17 @@ mooncake_master \ --tenant_quota_connector_uri=/etc/mooncake/tenant_quotas.yaml ``` -The v1 connector is a writable YAML file. The file must use schema version `1`; tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters; quotas must be positive integers with optional `B`, `KB`, `MB`, `GB`, or `TB` units: +You can also store the same YAML policy in etcd when Mooncake Store is built with `STORE_USE_ETCD=ON`: + +```bash +mooncake_master \ + --enable_multi_tenants=true \ + --cluster_id=mooncake_cluster \ + --tenant_quota_connector_type=etcd \ + --tenant_quota_connector_uri=127.0.0.1:2379 +``` + +The etcd connector stores the policy at `mooncake-store//tenant_quota_policy`. If the key does not exist, the master starts with an empty policy so the first tenant policy can be created through the admin API. It shares the process-wide store etcd client used by HA/oplog, so if HA or oplog also uses etcd, `tenant_quota_connector_uri` must match those etcd endpoints. The policy must use schema version `1`; tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters; quotas must be positive integers with optional `B`, `KB`, `MB`, `GB`, or `TB` units: ```yaml version: 1 @@ -475,8 +485,8 @@ mooncake_master \ | Flag | Default | Description | |------|---------|-------------| | `--enable_multi_tenants` | `false` | Enable strict tenant registration and per-tenant memory quota admission | -| `--tenant_quota_connector_type` | `file` | Tenant quota policy connector type | -| `--tenant_quota_connector_uri` | empty | Connector URI; for `file`, the writable YAML policy path | +| `--tenant_quota_connector_type` | `file` | Tenant quota policy connector type: `file` or `etcd` when built with `STORE_USE_ETCD=ON` | +| `--tenant_quota_connector_uri` | empty | Connector URI; for `file`, the writable YAML policy path; for `etcd`, the endpoints string | ### High Availability diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md index 4ebe9751af..72912c20c9 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -95,7 +95,7 @@ To reduce cache warm-up time after a master restart, the Master Service supports The Master Service can optionally enforce strict multi-tenant memory quota admission. This feature is disabled by default. When `enable_multi_tenants=false`, request tenant IDs are ignored for object placement, all objects use the `default` namespace, and tenant quota management requests return `UNAVAILABLE_IN_CURRENT_MODE`. -When strict multi-tenant mode is enabled, the tenant quota policy is loaded from the configured connector. The v1 connector is a writable YAML file configured by `tenant_quota_connector_type=file` and `tenant_quota_connector_uri=`. Tenants must be explicitly present in that connector policy before they can write. Missing tenants, empty tenants, and an unregistered `default` tenant are rejected with `TENANT_NOT_REGISTERED`. +When strict multi-tenant mode is enabled, the tenant quota policy is loaded from the configured connector. Supported connector types are `file` and, when the store is built with `STORE_USE_ETCD=ON`, `etcd`. The `file` connector uses `tenant_quota_connector_uri=` as a writable YAML policy path. The `etcd` connector uses `tenant_quota_connector_uri=` as the etcd endpoints string and stores the same YAML policy in `mooncake-store//tenant_quota_policy`; if that key does not exist, the master starts with an empty policy so the first policy can be created through the admin API. The etcd connector shares the process-wide store etcd client used by HA/oplog, so deployments that enable both must configure matching etcd endpoints. Tenants must be explicitly present in that connector policy before they can write. Missing tenants, empty tenants, and an unregistered `default` tenant are rejected with `TENANT_NOT_REGISTERED`. The YAML policy uses schema version `1`: diff --git a/mooncake-store/include/tenant_quota_policy_store.h b/mooncake-store/include/tenant_quota_policy_store.h index 30f3d85d41..ad6d2ae9c0 100644 --- a/mooncake-store/include/tenant_quota_policy_store.h +++ b/mooncake-store/include/tenant_quota_policy_store.h @@ -45,7 +45,24 @@ class YamlTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { std::mutex mutex_; }; +#ifdef STORE_USE_ETCD +class EtcdTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { + public: + EtcdTenantQuotaPolicyStore(const std::string& endpoints, + const std::string& cluster_id); + + tl::expected Load() override; + tl::expected Save( + const TenantQuotaPolicySnapshot& snapshot) override; + + private: + std::string key_; + std::mutex mutex_; +}; +#endif + tl::expected, std::string> -CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri); +CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri, + const std::string& cluster_id); } // namespace mooncake diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 5010c7053c..2259396b12 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -273,7 +273,8 @@ MasterService::MasterService(const MasterServiceConfig& config) if (enable_multi_tenants_) { auto store = CreateTenantQuotaPolicyStore(tenant_quota_connector_type_, - tenant_quota_connector_uri_); + tenant_quota_connector_uri_, + cluster_id_); if (!store) { throw std::invalid_argument(store.error()); } diff --git a/mooncake-store/src/tenant_quota_policy_store.cpp b/mooncake-store/src/tenant_quota_policy_store.cpp index 23e10b7099..8edef856ce 100644 --- a/mooncake-store/src/tenant_quota_policy_store.cpp +++ b/mooncake-store/src/tenant_quota_policy_store.cpp @@ -18,6 +18,9 @@ #include #include +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#endif #include "types.h" namespace mooncake { @@ -31,6 +34,34 @@ std::string ErrnoMessage(const std::string& action, const std::string& path) { return action + " '" + path + "' failed: " + std::strerror(errno); } +#ifdef STORE_USE_ETCD +tl::expected NormalizeClusterIdForEtcdKey( + const std::string& cluster_id) { + std::string normalized = + cluster_id.empty() ? std::string(DEFAULT_CLUSTER_ID) : cluster_id; + while (!normalized.empty() && normalized.back() == '/') { + normalized.pop_back(); + } + if (normalized.empty()) { + normalized = DEFAULT_CLUSTER_ID; + } + if (!IsValidClusterIdComponent(normalized)) { + return tl::make_unexpected("invalid tenant quota etcd cluster_id '" + + cluster_id + "'"); + } + return normalized; +} + +tl::expected BuildTenantQuotaEtcdKey( + const std::string& cluster_id) { + auto normalized = NormalizeClusterIdForEtcdKey(cluster_id); + if (!normalized) { + return tl::make_unexpected(normalized.error()); + } + return "mooncake-store/" + normalized.value() + "/tenant_quota_policy"; +} +#endif + tl::expected WriteAll(int fd, const std::string& content, const std::string& path) { const char* data = content.data(); @@ -345,17 +376,91 @@ tl::expected YamlTenantQuotaPolicyStore::Save( return {}; } +#ifdef STORE_USE_ETCD +EtcdTenantQuotaPolicyStore::EtcdTenantQuotaPolicyStore( + const std::string& endpoints, const std::string& cluster_id) { + auto key = BuildTenantQuotaEtcdKey(cluster_id); + if (!key) { + throw std::invalid_argument(key.error()); + } + if (endpoints.empty()) { + throw std::invalid_argument( + "tenant quota etcd connector requires a non-empty uri"); + } + ErrorCode connect_error = EtcdHelper::ConnectToEtcdStoreClient(endpoints); + if (connect_error != ErrorCode::OK) { + if (connect_error == ErrorCode::INVALID_PARAMS) { + throw std::runtime_error( + "failed to connect tenant quota etcd store: " + "tenant_quota_connector_uri must match the already connected " + "store etcd endpoints used by HA/oplog"); + } + throw std::runtime_error("failed to connect tenant quota etcd store: " + + toString(connect_error)); + } + key_ = std::move(key.value()); +} + +tl::expected +EtcdTenantQuotaPolicyStore::Load() { + std::lock_guard lock(mutex_); + std::string content; + EtcdRevisionId revision_id = 0; + ErrorCode error = + EtcdHelper::Get(key_.c_str(), key_.size(), content, revision_id); + if (error == ErrorCode::ETCD_KEY_NOT_EXIST) { + return TenantQuotaPolicySnapshot{}; + } + if (error != ErrorCode::OK) { + return tl::make_unexpected( + "failed to load tenant quota policy from " + "etcd key '" + + key_ + "': " + toString(error)); + } + return ParseTenantQuotaPolicyYaml(content); +} + +tl::expected EtcdTenantQuotaPolicyStore::Save( + const TenantQuotaPolicySnapshot& snapshot) { + std::lock_guard lock(mutex_); + const std::string content = FormatTenantQuotaPolicyYaml(snapshot); + ErrorCode error = EtcdHelper::Put(key_.c_str(), key_.size(), + content.c_str(), content.size()); + if (error != ErrorCode::OK) { + return tl::make_unexpected( + "failed to save tenant quota policy to " + "etcd key '" + + key_ + "': " + toString(error)); + } + return {}; +} +#endif + tl::expected, std::string> -CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri) { - if (type != "file") { - return tl::make_unexpected("unsupported tenant quota connector type '" + - type + "'"); +CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri, + const std::string& cluster_id) { + if (type == "file") { + if (uri.empty()) { + return tl::make_unexpected( + "tenant quota file connector requires a non-empty uri"); + } + return std::make_unique(uri); } - if (uri.empty()) { + if (type == "etcd") { +#ifdef STORE_USE_ETCD + try { + return std::make_unique(uri, + cluster_id); + } catch (const std::exception& e) { + return tl::make_unexpected(e.what()); + } +#else return tl::make_unexpected( - "tenant quota file connector requires a non-empty uri"); + "tenant quota etcd connector requires STORE_USE_ETCD"); +#endif } - return std::make_unique(uri); + return tl::make_unexpected("unsupported tenant quota connector type '" + + type + "'"); } } // namespace mooncake diff --git a/mooncake-store/tests/tenant_quota_test.cpp b/mooncake-store/tests/tenant_quota_test.cpp index 7faba55aff..c17f7257b3 100644 --- a/mooncake-store/tests/tenant_quota_test.cpp +++ b/mooncake-store/tests/tenant_quota_test.cpp @@ -2,11 +2,18 @@ #include "tenant_quota_policy_store.h" #include "types.h" +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#endif + +#include #include #include #include #include +#include #include +#include #include #include @@ -15,6 +22,19 @@ namespace mooncake { namespace { +#ifdef STORE_USE_ETCD +constexpr const char* kTenantQuotaEtcdEndpoints = "127.0.0.1:2379"; +constexpr std::string_view kTenantQuotaEtcdProbeKey = "tenant_quota_probe"; + +std::string GetTenantQuotaEtcdEndpoints() { + const char* endpoints = std::getenv("MOONCAKE_TENANT_QUOTA_ETCD_ENDPOINTS"); + if (endpoints != nullptr && endpoints[0] != '\0') { + return endpoints; + } + return kTenantQuotaEtcdEndpoints; +} +#endif + TenantQuotaSnapshot Snapshot(const TenantQuotaTable& table, const std::string& tenant_id) { auto snapshot = table.GetTenantSnapshot(tenant_id); @@ -36,6 +56,46 @@ std::filesystem::path MakeTempPolicyPath(const std::string& suffix) { std::to_string(::getpid()) + "_" + suffix + ".yaml"); } +#ifdef STORE_USE_ETCD +std::string PrefixEnd(std::string prefix) { + for (int i = static_cast(prefix.size()) - 1; i >= 0; --i) { + unsigned char c = static_cast(prefix[i]); + if (c < 0xFF) { + prefix[i] = static_cast(c + 1); + prefix.resize(i + 1); + return prefix; + } + } + return std::string(1, '\0'); +} + +std::optional GetTenantQuotaEtcdSkipReason() { + const std::string endpoints = GetTenantQuotaEtcdEndpoints(); + ErrorCode error = EtcdHelper::ConnectToEtcdStoreClient(endpoints); + if (error != ErrorCode::OK) { + return "Etcd server not reachable at " + endpoints + ": " + + toString(error); + } + std::string value; + EtcdRevisionId revision_id = 0; + error = + EtcdHelper::Get(kTenantQuotaEtcdProbeKey.data(), + kTenantQuotaEtcdProbeKey.size(), value, revision_id); + if (error == ErrorCode::ETCD_OPERATION_ERROR) { + return "Etcd server not reachable at " + endpoints + ": " + + toString(error); + } + return std::nullopt; +} + +void CleanupTenantQuotaEtcdCluster(const std::string& cluster_id) { + std::string prefix = "mooncake-store/" + cluster_id + "/"; + std::string end = PrefixEnd(prefix); + (void)EtcdHelper::DeleteRange(prefix.c_str(), prefix.size(), end.c_str(), + end.size()); +} +#endif + void MakeOrphanTenant(TenantQuotaTable* table, const std::string& tenant_id, uint64_t bytes) { ASSERT_TRUE(table->UpsertTenantPolicy(tenant_id, bytes).has_value()); @@ -228,6 +288,77 @@ TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlFile) { std::filesystem::remove(path); } +TEST(TenantQuotaPolicyStoreTest, FileFactoryCreatesYamlStore) { + const auto path = MakeTempPolicyPath("factory-file"); + auto store = + CreateTenantQuotaPolicyStore("file", path.string(), "test_cluster"); + ASSERT_TRUE(store.has_value()) << store.error(); +} + +TEST(TenantQuotaPolicyStoreTest, FileFactoryRequiresUri) { + auto store = CreateTenantQuotaPolicyStore("file", "", "test_cluster"); + ASSERT_FALSE(store.has_value()); + EXPECT_NE(store.error().find("non-empty uri"), std::string::npos); +} + +#ifndef STORE_USE_ETCD +TEST(TenantQuotaPolicyStoreTest, EtcdFactoryRequiresStoreUseEtcd) { + auto store = + CreateTenantQuotaPolicyStore("etcd", "127.0.0.1:2379", "test_cluster"); + ASSERT_FALSE(store.has_value()); + EXPECT_NE(store.error().find("STORE_USE_ETCD"), std::string::npos); +} +#endif + +#ifdef STORE_USE_ETCD +TEST(TenantQuotaPolicyStoreTest, EtcdMissingKeyLoadsEmptySnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } + + const std::string cluster_id = + "tenant_quota_missing_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); + + auto store = CreateTenantQuotaPolicyStore( + "etcd", GetTenantQuotaEtcdEndpoints(), cluster_id); + ASSERT_TRUE(store.has_value()) << store.error(); + + auto loaded = store.value()->Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_TRUE(loaded->tenant_quotas.empty()); + + CleanupTenantQuotaEtcdCluster(cluster_id); +} + +TEST(TenantQuotaPolicyStoreTest, EtcdRoundTripsSnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } + + const std::string cluster_id = + "tenant_quota_roundtrip_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); + + auto store = CreateTenantQuotaPolicyStore( + "etcd", GetTenantQuotaEtcdEndpoints(), cluster_id); + ASSERT_TRUE(store.has_value()) << store.error(); + + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"tenant-a", 1024}, {"tenant-b", 2048}}; + auto save = store.value()->Save(snapshot); + ASSERT_TRUE(save.has_value()) << save.error(); + + auto loaded = store.value()->Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_EQ(loaded->tenant_quotas, snapshot.tenant_quotas); + + CleanupTenantQuotaEtcdCluster(cluster_id); +} +#endif + TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlSpecialScalarNames) { TenantQuotaPolicySnapshot snapshot; snapshot.tenant_quotas = {{"foo#bar", 1}, From 023cb88f3077fd0f56d1b8982a0f3763b8dfc873 Mon Sep 17 00:00:00 2001 From: Schatten Date: Mon, 6 Jul 2026 10:42:02 +0800 Subject: [PATCH 12/57] [Store] Fix ConfigDict global segment size validation (#2661) Signed-off-by: Schatten --- mooncake-store/src/real_client.cpp | 18 +++++++++++++--- mooncake-store/tests/pybind_client_test.cpp | 24 ++++++++++++++++++--- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index dc544b7872..af18c8e1dd 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1025,7 +1025,17 @@ tl::expected RealClient::setup_internal( get_config(config, CONFIG_KEY_IPC_SOCKET_PATH); // A size of 0 keeps the pure client/server setup semantics. - auto validate_size = [](const char *key, size_t value) { + // global_segment_size is a total capacity and may exceed max_mr_size; the + // setup path splits it into mountable chunks below. + auto validate_min_size = [](const char *key, size_t value) { + if (value != 0 && value < MIN_SEGMENT_SIZE) { + LOG(ERROR) << "Invalid " << key << ": " << value + << ", must be 0 or at least " << MIN_SEGMENT_SIZE; + return false; + } + return true; + }; + auto validate_single_segment_size = [](const char *key, size_t value) { if ((value != 0 && value < MIN_SEGMENT_SIZE) || value > MAX_SEGMENT_SIZE) { LOG(ERROR) << "Invalid " << key << ": " << value @@ -1035,8 +1045,10 @@ tl::expected RealClient::setup_internal( } return true; }; - if (!validate_size(CONFIG_KEY_GLOBAL_SEGMENT_SIZE, global_segment_size) || - !validate_size(CONFIG_KEY_LOCAL_BUFFER_SIZE, local_buffer_size)) { + if (!validate_min_size(CONFIG_KEY_GLOBAL_SEGMENT_SIZE, + global_segment_size) || + !validate_single_segment_size(CONFIG_KEY_LOCAL_BUFFER_SIZE, + local_buffer_size)) { return tl::unexpected(ErrorCode::INVALID_PARAMS); } diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp index 27dc31c4a5..c08a48a8e7 100644 --- a/mooncake-store/tests/pybind_client_test.cpp +++ b/mooncake-store/tests/pybind_client_test.cpp @@ -992,6 +992,23 @@ TEST_F(RealClientTest, SetupWithConfigDictAllowsZeroSizes) { << "Setup should preserve zero-size pure client/server semantics"; } +TEST_F(RealClientTest, + ConfigDictGlobalSegmentSizeAboveMaxPassesSizeValidation) { + ConfigDict config = MakeConfigDict( + "localhost:17816", std::to_string(MAX_SEGMENT_SIZE + 1), "0"); + config[CONFIG_KEY_PROTOCOL] = "tcp"; + config[CONFIG_KEY_MASTER_SERVER_ADDR] = "127.0.0.1:1"; + + ::testing::internal::CaptureStderr(); + auto result = py_client_->setup_internal(config); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(logs.find("Invalid global_segment_size"), std::string::npos) + << "global_segment_size above MAX_SEGMENT_SIZE should be accepted as " + "total capacity by ConfigDict validation"; +} + TEST_F(RealClientTest, ErrSetupWithInvalidConfigDictSize) { GLogMuter muter; ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())) @@ -999,15 +1016,16 @@ TEST_F(RealClientTest, ErrSetupWithInvalidConfigDictSize) { master_address_ = master_.master_address(); struct InvalidSizeCase { - const char* local_hostname; - const char* global_segment_size; - const char* local_buffer_size; + std::string local_hostname; + std::string global_segment_size; + std::string local_buffer_size; }; const InvalidSizeCase invalid_size_cases[] = { {"localhost:17816", "50%", "16MB"}, {"localhost:17817", "16MB", "16XB"}, {"localhost:17818", "-5", "16MB"}, + {"localhost:17819", "0", std::to_string(MAX_SEGMENT_SIZE + 1)}, }; for (const auto& test_case : invalid_size_cases) { From 3adc3aaed2ca185906c368d9e6d0f3ac65fdd949 Mon Sep 17 00:00:00 2001 From: Aoi Date: Mon, 6 Jul 2026 11:00:31 +0800 Subject: [PATCH 13/57] [Store] Refactor store buffer headers and slice splitting (#2732) Co-authored-by: Aionw --- .../benchmarks/allocation_strategy_bench.cpp | 2 +- mooncake-store/benchmarks/allocator_bench.cpp | 2 +- ...ent_buffer.hpp => aligned_client_buffer.h} | 2 +- mooncake-store/include/allocator.h | 2 +- .../{client_buffer.hpp => client_buffer.h} | 2 +- mooncake-store/include/file_storage.h | 2 +- ...ffset_allocator.hpp => offset_allocator.h} | 2 +- mooncake-store/include/pyclient.h | 2 +- mooncake-store/include/real_client.h | 2 +- .../{serializer.hpp => serializer.h} | 2 +- mooncake-store/include/storage_backend.h | 4 +- mooncake-store/src/aligned_client_buffer.cpp | 2 +- mooncake-store/src/client_buffer.cpp | 2 +- mooncake-store/src/client_service.cpp | 2 +- mooncake-store/src/file_storage.cpp | 2 +- .../catalog_backed_snapshot_provider.cpp | 2 +- mooncake-store/src/master_service.cpp | 2 +- mooncake-store/src/offset_allocator.cpp | 4 +- mooncake-store/src/real_client.cpp | 69 +++---------------- mooncake-store/src/serialize/serializer.cpp | 4 +- mooncake-store/tests/client_buffer_test.cpp | 2 +- .../tests/client_local_hot_cache_test.cpp | 2 +- .../tests/ha/snapshot/snapshot_test_utils.h | 2 +- .../tests/offset_allocator_test.cpp | 2 +- mooncake-store/tests/serializer_test.cpp | 2 +- 25 files changed, 38 insertions(+), 85 deletions(-) rename mooncake-store/include/{aligned_client_buffer.hpp => aligned_client_buffer.h} (98%) rename mooncake-store/include/{client_buffer.hpp => client_buffer.h} (99%) rename mooncake-store/include/offset_allocator/{offset_allocator.hpp => offset_allocator.h} (99%) rename mooncake-store/include/serialize/{serializer.hpp => serializer.h} (99%) diff --git a/mooncake-store/benchmarks/allocation_strategy_bench.cpp b/mooncake-store/benchmarks/allocation_strategy_bench.cpp index 10e45fe7f7..adf8facb6f 100644 --- a/mooncake-store/benchmarks/allocation_strategy_bench.cpp +++ b/mooncake-store/benchmarks/allocation_strategy_bench.cpp @@ -15,7 +15,7 @@ #include #include "types.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "allocator.h" #include "allocation_strategy.h" diff --git a/mooncake-store/benchmarks/allocator_bench.cpp b/mooncake-store/benchmarks/allocator_bench.cpp index 8c4a8d908c..17cd58b2a7 100644 --- a/mooncake-store/benchmarks/allocator_bench.cpp +++ b/mooncake-store/benchmarks/allocator_bench.cpp @@ -6,7 +6,7 @@ #include #include -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" using namespace mooncake::offset_allocator; diff --git a/mooncake-store/include/aligned_client_buffer.hpp b/mooncake-store/include/aligned_client_buffer.h similarity index 98% rename from mooncake-store/include/aligned_client_buffer.hpp rename to mooncake-store/include/aligned_client_buffer.h index 436fc9bb92..1cfd44ecb7 100644 --- a/mooncake-store/include/aligned_client_buffer.hpp +++ b/mooncake-store/include/aligned_client_buffer.h @@ -1,6 +1,6 @@ #pragma once -#include "client_buffer.hpp" +#include "client_buffer.h" namespace mooncake { diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h index e28b7d5c54..de74e7ffea 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -7,7 +7,7 @@ #include #include "cachelib_memory_allocator/MemoryAllocator.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" using facebook::cachelib::MemoryAllocator; diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.h similarity index 99% rename from mooncake-store/include/client_buffer.hpp rename to mooncake-store/include/client_buffer.h index b9f84089fe..f3fce555fe 100644 --- a/mooncake-store/include/client_buffer.hpp +++ b/mooncake-store/include/client_buffer.h @@ -5,7 +5,7 @@ #include #include -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" #include "replica.h" diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 6aa2ac3d56..4a2f996310 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -1,7 +1,7 @@ #pragma once #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "storage_backend.h" #include "pinned_buffer_pool.h" diff --git a/mooncake-store/include/offset_allocator/offset_allocator.hpp b/mooncake-store/include/offset_allocator/offset_allocator.h similarity index 99% rename from mooncake-store/include/offset_allocator/offset_allocator.hpp rename to mooncake-store/include/offset_allocator/offset_allocator.h index 20052ee1ca..91da68c281 100644 --- a/mooncake-store/include/offset_allocator/offset_allocator.hpp +++ b/mooncake-store/include/offset_allocator/offset_allocator.h @@ -8,7 +8,7 @@ #include #include "mutex.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" namespace mooncake::offset_allocator { typedef unsigned char uint8; diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h index c2e2201c5d..0baefebf4f 100644 --- a/mooncake-store/include/pyclient.h +++ b/mooncake-store/include/pyclient.h @@ -11,7 +11,7 @@ #include #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "mutex.h" #include "utils.h" #include "file_storage.h" diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index feaa5c72ce..870fa9ea7e 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -14,7 +14,7 @@ #include "pyclient.h" #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "mutex.h" #include "utils.h" #include "rpc_types.h" diff --git a/mooncake-store/include/serialize/serializer.hpp b/mooncake-store/include/serialize/serializer.h similarity index 99% rename from mooncake-store/include/serialize/serializer.hpp rename to mooncake-store/include/serialize/serializer.h index cb153b1a84..135e345ee4 100644 --- a/mooncake-store/include/serialize/serializer.hpp +++ b/mooncake-store/include/serialize/serializer.h @@ -141,4 +141,4 @@ class SerializationHelper { } }; -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 5fd071b59b..8a33264acc 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -13,7 +13,7 @@ #include "file_interface.h" #include "mutex.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" namespace mooncake { @@ -1209,4 +1209,4 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { tl::expected, ErrorCode> CreateStorageBackend(const FileStorageConfig& config); -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/src/aligned_client_buffer.cpp b/mooncake-store/src/aligned_client_buffer.cpp index 55d83e5a1d..7569868e42 100644 --- a/mooncake-store/src/aligned_client_buffer.cpp +++ b/mooncake-store/src/aligned_client_buffer.cpp @@ -1,4 +1,4 @@ -#include "aligned_client_buffer.hpp" +#include "aligned_client_buffer.h" #include #include diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp index db102c3b16..6aaea4af10 100644 --- a/mooncake-store/src/client_buffer.cpp +++ b/mooncake-store/src/client_buffer.cpp @@ -1,4 +1,4 @@ -#include "client_buffer.hpp" +#include "client_buffer.h" #include #include diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 7c078ad7d9..446503c91a 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -35,7 +35,7 @@ #include "config.h" #include "ha/leadership/leader_coordinator_factory.h" #include "types.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "utils.h" #include "rpc_types.h" #include "local_hot_cache.h" diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 00e7db4bcd..58bb601aaa 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -4,7 +4,7 @@ #include #include -#include "aligned_client_buffer.hpp" +#include "aligned_client_buffer.h" #include "storage_backend.h" #include "client_metric.h" #include "utils.h" diff --git a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp index dc150e8ea9..dc5c8f6fc1 100644 --- a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp +++ b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp @@ -16,7 +16,7 @@ #include "ha/snapshot/catalog/snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "segment.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "utils/zstd_util.h" namespace mooncake { diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 2259396b12..5f55c2ed76 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -39,7 +39,7 @@ #include "ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "types.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "ha/snapshot/snapshot_logger.h" #include "utils/zstd_util.h" #include "utils/file_util.h" diff --git a/mooncake-store/src/offset_allocator.cpp b/mooncake-store/src/offset_allocator.cpp index 68b44025f6..60b2de3d9c 100644 --- a/mooncake-store/src/offset_allocator.cpp +++ b/mooncake-store/src/offset_allocator.cpp @@ -1,7 +1,7 @@ // (C) Sebastian Aaltonen 2023 // MIT License (see file: LICENSE) -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include #include @@ -683,4 +683,4 @@ std::ostream& operator<<(std::ostream& os, return os; } -} // namespace mooncake::offset_allocator \ No newline at end of file +} // namespace mooncake::offset_allocator diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index af18c8e1dd..2f35506e5a 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -18,7 +18,7 @@ #include #include "real_client.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "common.h" #include "config.h" #include "mutex.h" @@ -2625,7 +2625,7 @@ std::shared_ptr RealClient::get_buffer_internal( } // MEMORY / DISK: use client_->Get. FilterQueryResult ensures - // Client::Get's internal FindFirstCompleteReplica can only see + // Client::Get internal FindFirstCompleteReplica can only see // the replica we selected, preventing accidental LOCAL_DISK picks. auto runtime_accelerator = device::GetAcceleratorRegistry().RuntimeAccelerators(); @@ -3647,17 +3647,7 @@ std::vector> RealClient::batch_put_from_internal( void *buffer = buffers[i]; size_t size = sizes[i]; - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - - all_slices[key] = std::move(slices); + all_slices[key] = split_into_slices(buffer, size); } std::vector> ordered_batched_slices; @@ -3697,15 +3687,7 @@ tl::expected RealClient::put_from_internal( } // Create slices directly from the user buffer - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = split_into_slices(buffer, size); auto put_result = client_->Put(key, slices, config); if (!put_result) { @@ -3812,14 +3794,7 @@ tl::expected RealClient::upsert_from_internal( return {}; } - std::vector slices; - uint64_t offset = 0; - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = split_into_slices(buffer, size); auto result = client_->Upsert(key, slices, config); if (!result) { @@ -3865,15 +3840,8 @@ RealClient::batch_upsert_from_internal(const std::vector &keys, ordered_batched_slices.reserve(keys.size()); for (size_t i = 0; i < keys.size(); ++i) { - std::vector slices; - uint64_t offset = 0; - while (offset < sizes[i]) { - auto chunk_size = std::min(sizes[i] - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffers[i]) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - ordered_batched_slices.emplace_back(std::move(slices)); + ordered_batched_slices.emplace_back( + split_into_slices(buffers[i], sizes[i])); } return client_->BatchUpsert(keys, ordered_batched_slices, config); @@ -4750,25 +4718,10 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer, } // Create slices directly from the user buffer - std::vector slices; - // Add metadata slice - uint64_t metadata_offset = 0; - while (metadata_offset < metadata_size) { - auto metadata_chunk_size = - std::min(metadata_size - metadata_offset, kMaxSliceSize); - void *metadata_chunk_ptr = - static_cast(metadata_buffer) + metadata_offset; - slices.emplace_back(Slice{metadata_chunk_ptr, metadata_chunk_size}); - metadata_offset += metadata_chunk_size; - } - - uint64_t offset = 0; - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = + split_into_slices(metadata_buffer, metadata_size); + auto data_slices = split_into_slices(buffer, size); + slices.insert(slices.end(), data_slices.begin(), data_slices.end()); auto put_result = client_->Put(key, slices, config); if (!put_result) { LOG(ERROR) << "Put operation failed with error: " diff --git a/mooncake-store/src/serialize/serializer.cpp b/mooncake-store/src/serialize/serializer.cpp index 00b46d47d5..acfc5feb01 100644 --- a/mooncake-store/src/serialize/serializer.cpp +++ b/mooncake-store/src/serialize/serializer.cpp @@ -1,8 +1,8 @@ #include #include -#include "serialize/serializer.hpp" -#include "offset_allocator/offset_allocator.hpp" +#include "serialize/serializer.h" +#include "offset_allocator/offset_allocator.h" #include "types.h" #include "master_service.h" #include "utils/zstd_util.h" diff --git a/mooncake-store/tests/client_buffer_test.cpp b/mooncake-store/tests/client_buffer_test.cpp index e48069eca3..8a091f9b5a 100644 --- a/mooncake-store/tests/client_buffer_test.cpp +++ b/mooncake-store/tests/client_buffer_test.cpp @@ -1,5 +1,5 @@ // client_buffer_test.cpp -#include "client_buffer.hpp" +#include "client_buffer.h" #include #include diff --git a/mooncake-store/tests/client_local_hot_cache_test.cpp b/mooncake-store/tests/client_local_hot_cache_test.cpp index 86927c2814..70fbf7c162 100644 --- a/mooncake-store/tests/client_local_hot_cache_test.cpp +++ b/mooncake-store/tests/client_local_hot_cache_test.cpp @@ -1,6 +1,6 @@ // client_local_hot_cache_test.cpp #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "count_min_sketch.h" #include "local_hot_cache.h" #include "replica.h" diff --git a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h index a160c1433a..01262f671c 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h +++ b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h @@ -19,7 +19,7 @@ #include "master_config.h" #include "replica.h" #include "segment.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "types.h" #include "utils/zstd_util.h" diff --git a/mooncake-store/tests/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index d0e5519e44..b544798e99 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -1,4 +1,4 @@ -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "mutex.h" #include "serializer.h" #include "types.h" diff --git a/mooncake-store/tests/serializer_test.cpp b/mooncake-store/tests/serializer_test.cpp index 5d9d6c7318..538d2ee05c 100644 --- a/mooncake-store/tests/serializer_test.cpp +++ b/mooncake-store/tests/serializer_test.cpp @@ -2,7 +2,7 @@ #include #include "serializer.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "segment.h" namespace mooncake::test { From d4a1f3961bde52020c503146ff2482ff59e8f0d0 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Mon, 6 Jul 2026 15:19:27 +0800 Subject: [PATCH 14/57] [TENT] SelectionPolicy: per-policy SL/TC/qp_pool schema (RFC #2568 step 1) (#2640) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] SelectionPolicy: add per-policy SL/TC/qp_pool schema (RFC #2568 step 1) Step 1 of the two-step plan agreed with @staryxchen on #2568: extend the existing `SelectionPolicy` (the opaque `policy_name` hook, not a new enum) with link-layer QoS attributes, so fabric QoS can be configured per policy instead of only process-globally via MC_IB_SL / MC_IB_TC. This is schema + plumbing only. It does NOT yet apply SL/TC at QP setup — real per-class differentiation needs per-class QP pools, which is the explicit step-2 follow-up. `qp_pool` is parsed and reserved here so the JSON schema is forward-compatible and step 2 needs no schema change or breakage of existing configs (per @staryxchen's reminder). - SelectionPolicy: add optional `service_level` (0-15), `traffic_class` (0-255), `qp_pool` (string). nullopt = fall back to the global RdmaParams default = today's behavior. - loadPolicies(): parse the three from the same JSON as existing fields; out-of-range SL/TC are ignored with a warning so a bad config never changes selection. - SelectionResult: carry the matched policy's SL/TC/qp_pool out to the caller for the step-2 QP-setup follow-up. - Tests: parse + carry-through, and out-of-range-ignored. Fully backward compatible: unset fields behave exactly as before, and the default policies (which don't set them) are unaffected. Note: not built locally (no TENT toolchain on my dev machine); relying on CI. clang-format (v20.1.8) applied. * [TENT] Address review: qp_pool empty/non-string handling + trim comments Per @staryxchen's review on #2640: - Treat an empty-string qp_pool the same as unset (nullopt = default pool), so a blank config value doesn't look like an explicit pool. - Warn on a non-string qp_pool (e.g. `"qp_pool": 42`) instead of silently no-op'ing it, matching the SL/TC out-of-range warnings — a typo'd config should not look like it took effect. - Trim the RFC-reference / step-split comments in the header. - Add a test covering empty / non-string qp_pool -> unset. --------- Co-authored-by: catyans --- .../include/tent/runtime/transport_selector.h | 12 ++ .../tent/src/runtime/transport_selector.cpp | 43 +++++++ .../tent/tests/transport_selector_test.cpp | 109 ++++++++++++++++++ 3 files changed, 164 insertions(+) diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h index eaa2341f33..a419e58d69 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h @@ -117,6 +117,14 @@ struct SelectionPolicy { // Transport preference list (evaluated in order) std::vector transports; + + // Per-policy link-layer QoS. nullopt = fall back to the global RdmaParams + // value. InfiniBand Service Level (0-15) and Traffic Class / DSCP (0-255). + std::optional service_level; + std::optional traffic_class; + // Named QP pool this policy's traffic should land on; parsed and stored for + // now, routing to be wired later. Unset = the current single "data QP". + std::optional qp_pool; }; /** @@ -125,6 +133,10 @@ struct SelectionPolicy { struct SelectionResult { TransportType transport = UNSPEC; uint64_t device_mask = ~0ULL; // Bitmask of allowed devices (~0 = all) + // Resolved link-layer QoS from the matched policy (nullopt = default). + std::optional service_level; + std::optional traffic_class; + std::optional qp_pool; }; /** diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp index fc53df32cf..046af38c2e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp @@ -218,6 +218,42 @@ void TransportSelector::loadPolicies() { } } + // Parse link-layer QoS attributes (RFC #2519 / #2568, step 1: stored + // only, not yet applied to QPs). Out-of-range values are ignored so a + // bad config never breaks selection. + if (policy_json.contains("service_level")) { + int sl = policy_json.value("service_level", -1); + if (sl >= 0 && sl <= 15) { + policy.service_level = sl; + } else { + LOG(WARNING) << "Ignore service_level in policy " << policy.name + << ", value " << sl << " out of range (0-15)"; + } + } + if (policy_json.contains("traffic_class")) { + int tc = policy_json.value("traffic_class", -1); + if (tc >= 0 && tc <= 255) { + policy.traffic_class = tc; + } else { + LOG(WARNING) << "Ignore traffic_class in policy " << policy.name + << ", value " << tc << " out of range (0-255)"; + } + } + // Reserved for step 2 (per-class QP pools); parsed for forward schema + // compatibility, no effect yet. + if (policy_json.contains("qp_pool")) { + auto& qp = policy_json["qp_pool"]; + if (!qp.is_string()) { + LOG(WARNING) << "Ignore qp_pool in policy " << policy.name + << ", expected a string"; + } else { + auto value = qp.get(); + // Treat an empty string the same as unset (use default pool) + // so a blank config value doesn't look like an explicit pool. + if (!value.empty()) policy.qp_pool = std::move(value); + } + } + policies_.push_back(std::move(policy)); LOG(INFO) << "Loaded transport policy: " << policy.name << " (segment_type=" << segment_type_str @@ -405,6 +441,13 @@ SelectionResult TransportSelector::select( return result; // UNSPEC, all devices } + // Carry the matched policy's link-layer QoS out to the caller (RFC #2519 / + // #2568, step 1). These are plumbed but not yet applied at QP setup; that + // is the per-class QP pool follow-up (step 2). + result.service_level = matching_policy->service_level; + result.traffic_class = matching_policy->traffic_class; + result.qp_pool = matching_policy->qp_pool; + // Convert device names to mask result.device_mask = ~0ULL; // Default: all devices if (!matching_policy->devices.empty() && topology_) { diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index a479c8e43a..4fb0f39018 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -641,6 +641,115 @@ TEST(TransportSelectorTest, HintNotInMatchingPolicyReturnsUnspec) { EXPECT_EQ(r.transport, UNSPEC); } +// RFC #2519 / #2568 step 1: a policy's link-layer QoS (service_level / +// traffic_class / qp_pool) is parsed from JSON and carried out via +// SelectionResult. (Step 1 only plumbs the values; applying them at QP setup +// is the per-class QP pool follow-up.) +TEST(TransportSelectorTest, PolicyLinkLayerQoSIsParsedAndCarried) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "kv-critical"; + policy["segment_type"] = "memory"; + policy["transports"] = {"rdma"}; + policy["service_level"] = 3; + policy["traffic_class"] = 96; + policy["qp_pool"] = "kv"; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "kv-critical"; + + auto r = selector.select(ctx, transports, /*index=*/0); + ASSERT_TRUE(r.service_level.has_value()); + EXPECT_EQ(r.service_level.value(), 3); + ASSERT_TRUE(r.traffic_class.has_value()); + EXPECT_EQ(r.traffic_class.value(), 96); + ASSERT_TRUE(r.qp_pool.has_value()); + EXPECT_EQ(r.qp_pool.value(), "kv"); +} + +// Out-of-range SL/TC are ignored (left as nullopt) so a bad config never +// changes selection behavior. +TEST(TransportSelectorTest, PolicyLinkLayerQoSOutOfRangeIgnored) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "bad-qos"; + policy["segment_type"] = "memory"; + policy["transports"] = {"rdma"}; + policy["service_level"] = 99; // > 15, invalid + policy["traffic_class"] = 9999; // > 255, invalid + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "bad-qos"; + + auto r = selector.select(ctx, transports, /*index=*/0); + EXPECT_FALSE(r.service_level.has_value()); + EXPECT_FALSE(r.traffic_class.has_value()); +} + +// An empty-string or non-string qp_pool is treated as unset (nullopt), so a +// blank / mistyped config value never looks like an explicit pool. +TEST(TransportSelectorTest, PolicyQpPoolEmptyOrNonStringIsUnset) { + auto conf = std::make_shared(); + json empty_pool; + empty_pool["name"] = "empty-pool"; + empty_pool["segment_type"] = "memory"; + empty_pool["transports"] = {"rdma"}; + empty_pool["qp_pool"] = ""; // empty -> unset + json bad_pool; + bad_pool["name"] = "bad-pool"; + bad_pool["segment_type"] = "memory"; + bad_pool["transports"] = {"rdma"}; + bad_pool["qp_pool"] = 42; // non-string -> ignored + conf->set("policy", json::array({empty_pool, bad_pool})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + + ctx.policy_name = "empty-pool"; + EXPECT_FALSE( + selector.select(ctx, transports, /*index=*/0).qp_pool.has_value()); + ctx.policy_name = "bad-pool"; + EXPECT_FALSE( + selector.select(ctx, transports, /*index=*/0).qp_pool.has_value()); +} + } // namespace } // namespace tent } // namespace mooncake From 29b22aa6ec11fca90c51e968c640adf32344b1bd Mon Sep 17 00:00:00 2001 From: LZW <99333079+Lin-z-w@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:06:26 +0800 Subject: [PATCH 15/57] [Store] feat: allow configuring client tenant id (#2755) --- .../mooncake-store-deployment-guide.md | 9 +- mooncake-integration/store/store_py.cpp | 3 +- mooncake-store/src/real_client_main.cpp | 4 +- mooncake-wheel/mooncake/mooncake_config.py | 14 +++- .../mooncake/mooncake_store_service.py | 3 +- mooncake-wheel/tests/test_mooncake_config.py | 84 ++++++++++++++++++- .../tests/test_mooncake_store_service_api.py | 65 ++++++++++++++ 7 files changed, 173 insertions(+), 9 deletions(-) diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 14be116665..9ad0f3aebd 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -675,6 +675,7 @@ The store service CLI only accepts `--config`, `-D/--define`, `--port`, and `--m | `MOONCAKE_LOCAL_HOSTNAME` | `local_hostname` | `localhost` | | | `MOONCAKE_OFFLOAD_ENABLED` | `enable_ssd_offload` | `false` | Client-side SSD offload | | `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `ssd_offload_path` | empty | Offload directory | +| `MOONCAKE_TENANT_ID` | `tenant_id` | `default` | Tenant identifier | | `MOONCAKE_CONFIG_PATH` | — | unset | Path to a JSON config file (takes precedence over the variables above) | ```{note} @@ -705,12 +706,14 @@ Or via a JSON config file. The service also exposes a lightweight HTTP API (on ` "local_buffer_size": 268435456, "protocol": "tcp", "device_name": "", - "master_server_address": "127.0.0.1:50051" + "master_server_address": "127.0.0.1:50051", + "tenant_id": "default" } ``` ```bash python -m mooncake.mooncake_store_service --config= --port=8081 +python -m mooncake.mooncake_store_service --config= -Dtenant_id=tenant-a ``` ### Method C — Resource-owning Real Client (`mooncake_client`) @@ -721,7 +724,8 @@ Run the `mooncake_client` binary as a standalone RPC process that owns storage r mooncake_client \ --global_segment_size="4GB" \ --master_server_address="127.0.0.1:50051" \ - --metadata_server="http://127.0.0.1:8080/metadata" + --metadata_server="http://127.0.0.1:8080/metadata" \ + --tenant_id="default" ``` | Flag | Default | Description | @@ -734,6 +738,7 @@ mooncake_client \ | `--protocol` | `tcp` | Transfer protocol | | `--device_names` | empty | Transfer device name(s), comma-separated | | `--threads` | `1` | Client worker thread count | +| `--tenant_id` | `default` | Tenant identifier | | `--enable_offload` | `false` | Enable client-side SSD offload | | `--start_offload_rpc_server` | `true` | Start the offload RPC server for dummy clients | diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index ec97258201..15ca9bac56 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -2167,7 +2167,8 @@ PYBIND11_MODULE(store, m) { " ipc_socket_path: IPC socket path.\n" " enable_ssd_offload: Enable SSD offload (default false).\n" " ssd_offload_path: SSD storage directory path (overrides env " - "var).") + "var).\n" + " tenant_id: Tenant identifier (default 'default').") .def( "setup_dummy", [](MooncakeStorePyWrapper &self, size_t mem_pool_size, diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 15c7202f33..fc1528b063 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -18,6 +18,7 @@ DEFINE_string(protocol, "tcp", "Protocol"); DEFINE_int32(port, 50052, "Real Client service port"); DEFINE_string(global_segment_size, "4 GB", "Size of global segment"); DEFINE_int32(threads, 1, "Number of threads for client service"); +DEFINE_string(tenant_id, "default", "Tenant identifier"); DEFINE_bool(enable_offload, false, "Enable offload availability"); DEFINE_bool(start_offload_rpc_server, true, "Expose TCP RPC for disk-tier reads " @@ -112,7 +113,8 @@ int main(int argc, char *argv[]) { FLAGS_host, FLAGS_metadata_server, global_segment_size, 0, FLAGS_protocol, FLAGS_device_names, FLAGS_master_server_address, nullptr, "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock", - FLAGS_port, FLAGS_enable_offload, FLAGS_start_offload_rpc_server); + FLAGS_port, FLAGS_enable_offload, FLAGS_start_offload_rpc_server, "", + FLAGS_tenant_id); if (!res) { LOG(FATAL) << "Failed to setup client: " << toString(res.error()); return -1; diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/mooncake-wheel/mooncake/mooncake_config.py index e311296865..8a7b303756 100644 --- a/mooncake-wheel/mooncake/mooncake_config.py +++ b/mooncake-wheel/mooncake/mooncake_config.py @@ -129,6 +129,7 @@ class MooncakeConfig: master_server_address (str): The address of the master server. enable_ssd_offload (bool): Enable SSD offload. Default is False. ssd_offload_path (str): The path to the SSD directory for offloading. + tenant_id (str): Tenant identifier. Default is "default". Example of configuration file: { @@ -140,7 +141,8 @@ class MooncakeConfig: "device_name": "", "master_server_address": "localhost:8081", "enable_ssd_offload": true, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "default" } For RDMA: @@ -153,7 +155,8 @@ class MooncakeConfig: "device_name": "mlx5_0", "master_server_address": "master:8081", "enable_ssd_offload": true, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "default" } """ local_hostname: str @@ -165,6 +168,7 @@ class MooncakeConfig: master_server_address: str enable_ssd_offload: bool = False ssd_offload_path: str = "" + tenant_id: str = "default" @staticmethod def from_file(file_path: str) -> 'MooncakeConfig': @@ -179,6 +183,8 @@ def from_file(file_path: str) -> 'MooncakeConfig': for field in required_fields: if field not in config: raise ValueError(f"Missing required config field: {field}") + ssd_offload_path = config.get("ssd_offload_path") + tenant_id = config.get("tenant_id") return MooncakeConfig( local_hostname=config.get("local_hostname"), metadata_server=config.get("metadata_server"), @@ -192,7 +198,8 @@ def from_file(file_path: str) -> 'MooncakeConfig': device_name=config.get("device_name", ""), master_server_address=config.get("master_server_address"), enable_ssd_offload=_parse_bool(config.get("enable_ssd_offload", False)), - ssd_offload_path=str(config.get("ssd_offload_path", "")), + ssd_offload_path=str(ssd_offload_path) if ssd_offload_path is not None else "", + tenant_id=str(tenant_id) if tenant_id is not None else "default", ) @staticmethod @@ -221,5 +228,6 @@ def load_from_env() -> 'MooncakeConfig': master_server_address=os.getenv("MOONCAKE_MASTER"), enable_ssd_offload=_parse_bool(os.getenv("MOONCAKE_OFFLOAD_ENABLED", "false")), ssd_offload_path=os.getenv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", ""), + tenant_id=os.getenv("MOONCAKE_TENANT_ID", "default"), ) return MooncakeConfig.from_file(config_file_path) \ No newline at end of file diff --git a/mooncake-wheel/mooncake/mooncake_store_service.py b/mooncake-wheel/mooncake/mooncake_store_service.py index a91584070d..3b3117576e 100644 --- a/mooncake-wheel/mooncake/mooncake_store_service.py +++ b/mooncake-wheel/mooncake/mooncake_store_service.py @@ -131,7 +131,8 @@ async def start_store_service(self, max_wait_time: float = 60): self.config.master_server_address, None, self.config.enable_ssd_offload, - self.config.ssd_offload_path + self.config.ssd_offload_path, + self.config.tenant_id ) if ret != 0: diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/mooncake-wheel/tests/test_mooncake_config.py index 26ead11351..5172682205 100644 --- a/mooncake-wheel/tests/test_mooncake_config.py +++ b/mooncake-wheel/tests/test_mooncake_config.py @@ -27,7 +27,8 @@ def setUp(self): "protocol": "tcp", "device_name": "eth0", "enable_ssd_offload": True, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "tenant-a" } def tearDown(self): @@ -52,6 +53,7 @@ def test_load_valid_config(self): self.assertEqual(config.device_name, "eth0") self.assertEqual(config.enable_ssd_offload, True) self.assertEqual(config.ssd_offload_path, "/nvme/mooncake_offload") + self.assertEqual(config.tenant_id, "tenant-a") def test_load_with_default_values(self): """Test loading configuration with default values""" @@ -69,6 +71,40 @@ def test_load_with_default_values(self): self.assertEqual(config.device_name, "") self.assertEqual(config.enable_ssd_offload, False) self.assertEqual(config.ssd_offload_path, "") + self.assertEqual(config.tenant_id, "default") + + def test_load_tenant_id_from_file(self): + """Test loading tenant_id from configuration file""" + self.write_config({**self.valid_config, "tenant_id": "tenant-from-file"}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "tenant-from-file") + + def test_tenant_id_defaults(self): + """Test tenant_id defaults to default when omitted""" + minimal_config = { + "local_hostname": "localhost", + "metadata_server": "localhost:8080", + "master_server_address": "localhost:8081" + } + self.write_config(minimal_config) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "default") + + def test_tenant_id_null_defaults(self): + """Test tenant_id defaults to default when explicitly null""" + self.write_config({**self.valid_config, "tenant_id": None}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "default") + + def test_ssd_offload_path_null_defaults(self): + """Test ssd_offload_path defaults to empty when explicitly null""" + self.write_config({**self.valid_config, "ssd_offload_path": None}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.ssd_offload_path, "") def test_enable_ssd_offload_string_values(self): """from_file must parse string booleans like load_from_env, and reject typos. @@ -144,6 +180,7 @@ def test_load_from_config_env(self): os.environ['MOONCAKE_DEVICE'] = self.valid_config["device_name"] os.environ['MOONCAKE_OFFLOAD_ENABLED'] = str(self.valid_config["enable_ssd_offload"]) os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH'] = self.valid_config["ssd_offload_path"] + os.environ['MOONCAKE_TENANT_ID'] = self.valid_config["tenant_id"] try: config = MooncakeConfig.load_from_env() @@ -155,6 +192,7 @@ def test_load_from_config_env(self): self.assertEqual(config.device_name, self.valid_config["device_name"]) self.assertEqual(config.enable_ssd_offload, self.valid_config["enable_ssd_offload"]) self.assertEqual(config.ssd_offload_path, self.valid_config["ssd_offload_path"]) + self.assertEqual(config.tenant_id, self.valid_config["tenant_id"]) finally: # Clean up environment variable @@ -166,6 +204,50 @@ def test_load_from_config_env(self): del os.environ['MOONCAKE_DEVICE'] del os.environ['MOONCAKE_OFFLOAD_ENABLED'] del os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH'] + del os.environ['MOONCAKE_TENANT_ID'] + + def test_tenant_id_from_env(self): + """Test loading tenant_id from MOONCAKE_TENANT_ID""" + previous_config_path = os.environ.pop("MOONCAKE_CONFIG_PATH", None) + previous_master = os.environ.pop("MOONCAKE_MASTER", None) + previous_tenant_id = os.environ.pop("MOONCAKE_TENANT_ID", None) + + os.environ["MOONCAKE_MASTER"] = self.valid_config["master_server_address"] + os.environ["MOONCAKE_TENANT_ID"] = "tenant-from-env" + + try: + config = MooncakeConfig.load_from_env() + self.assertEqual(config.tenant_id, "tenant-from-env") + finally: + os.environ.pop("MOONCAKE_MASTER", None) + os.environ.pop("MOONCAKE_TENANT_ID", None) + if previous_config_path is not None: + os.environ["MOONCAKE_CONFIG_PATH"] = previous_config_path + if previous_master is not None: + os.environ["MOONCAKE_MASTER"] = previous_master + if previous_tenant_id is not None: + os.environ["MOONCAKE_TENANT_ID"] = previous_tenant_id + + def test_tenant_id_env_defaults(self): + """Test tenant_id defaults to default when MOONCAKE_TENANT_ID is omitted""" + previous_config_path = os.environ.pop("MOONCAKE_CONFIG_PATH", None) + previous_master = os.environ.pop("MOONCAKE_MASTER", None) + previous_tenant_id = os.environ.pop("MOONCAKE_TENANT_ID", None) + + os.environ["MOONCAKE_MASTER"] = self.valid_config["master_server_address"] + + try: + config = MooncakeConfig.load_from_env() + self.assertEqual(config.tenant_id, "default") + finally: + os.environ.pop("MOONCAKE_MASTER", None) + os.environ.pop("MOONCAKE_TENANT_ID", None) + if previous_config_path is not None: + os.environ["MOONCAKE_CONFIG_PATH"] = previous_config_path + if previous_master is not None: + os.environ["MOONCAKE_MASTER"] = previous_master + if previous_tenant_id is not None: + os.environ["MOONCAKE_TENANT_ID"] = previous_tenant_id def test_load_from_env_missing(self): """Test loading configuration from environment variable when not set""" diff --git a/mooncake-wheel/tests/test_mooncake_store_service_api.py b/mooncake-wheel/tests/test_mooncake_store_service_api.py index 62d28e0e42..a2c07e5eb2 100644 --- a/mooncake-wheel/tests/test_mooncake_store_service_api.py +++ b/mooncake-wheel/tests/test_mooncake_store_service_api.py @@ -2,10 +2,12 @@ import asyncio import json import sys +import tempfile import types import unittest from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -50,6 +52,11 @@ def __init__(self): self.unmount_failures = set() self.allocated_mount_calls = [] self.free_unmount_calls = [] + self.setup_calls = [] + + def setup(self, *args): + self.setup_calls.append(args) + return 0 def mount_segment(self, path, size, offset, protocol, location): self.mount_calls.append((path, size, offset, protocol, location)) @@ -107,6 +114,64 @@ async def asyncSetUp(self): self.service.last_mount_info = {} self.service._state_lock = asyncio.Lock() + async def test_start_store_service_passes_tenant_id_to_setup(self): + fake_store = FakeStore() + self.service.config = SimpleNamespace( + local_hostname="localhost", + metadata_server="P2PHANDSHAKE", + global_segment_size=1024, + local_buffer_size=2048, + protocol="tcp", + device_name="", + master_server_address="127.0.0.1:50051", + enable_ssd_offload=False, + ssd_offload_path="", + tenant_id="tenant-a", + ) + + with patch( + "mooncake.mooncake_store_service.MooncakeDistributedStore", + return_value=fake_store, + ): + result = await self.service.start_store_service(max_wait_time=1) + + self.assertTrue(result) + self.assertEqual( + fake_store.setup_calls, + [ + ( + "localhost", + "P2PHANDSHAKE", + 1024, + 2048, + "tcp", + "", + "127.0.0.1:50051", + None, + False, + "", + "tenant-a", + ) + ], + ) + + async def test_cli_config_can_override_tenant_id(self): + config = { + "local_hostname": "localhost", + "metadata_server": "P2PHANDSHAKE", + "master_server_address": "127.0.0.1:50051", + "tenant_id": "tenant-from-file", + } + + with tempfile.TemporaryDirectory() as tmpdir: + config_path = Path(tmpdir) / "config.json" + config_path.write_text(json.dumps(config)) + service = MooncakeStoreService( + str(config_path), {"tenant_id": "tenant-from-cli"} + ) + + self.assertEqual(service.config.tenant_id, "tenant-from-cli") + async def test_mount_shm_then_unmount_shm_api(self): mount_resp = await self.service.handle_mount_shm( FakeRequest( From 118382f48b35d7bb53febdaa23dbfcddb725a05a Mon Sep 17 00:00:00 2001 From: xiangui <120565419+xiangui33423@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:27:44 +0800 Subject: [PATCH 16/57] [TE] Fix cross-GPU NVLink/MNNVL event-stream device mismatch (#2722) (#2754) * [TENT] Fix cross-GPU NVLink/MNNVL event-stream device mismatch (#2722) Fix 'cudaEventRecord: invalid resource handle' error when NVLink/MNNVL transfers target a GPU different from the current CUDA device. Root cause: CUDAStreamPool creates streams on the target device but restores the caller's original device afterward. Later, cudaEventCreate and cudaEventRecord execute on the caller's device, creating a device mismatch between the event and the stream. Solution: - Store the stream's device ID in NVLinkSubBatch/MnnvlSubBatch - Wrap event creation and recording in a device guard that switches to the stream's device - Restore the original device on all paths (success and error) This fix ensures events and streams are created on the same device, resolving cross-GPU transfer failures while maintaining correct device context for the caller. Co-Authored-By: Claude Opus 4.8 (1M context) * Fix CHECK_CUDA macro usage in void functions Replace CHECK_CUDA with manual error handling in startTransfer functions to fix compilation errors. CHECK_CUDA is designed to return Status objects, but startTransfer has void return type. Changes: - Manually check cudaGetDevice/cudaSetDevice errors in device switching - Log errors and mark tasks as FAILED on failure - Call cudaSetDevice directly without CHECK_CUDA when restoring device Co-Authored-By: Claude Opus 4.8 (1M context) * Apply clang-format to fix code formatting violations Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../tent/transport/mnnvl/mnnvl_transport.h | 1 + .../tent/transport/nvlink/nvlink_transport.h | 1 + .../src/transport/mnnvl/mnnvl_transport.cpp | 39 +++++++++++++++++++ .../src/transport/nvlink/nvlink_transport.cpp | 39 +++++++++++++++++++ 4 files changed, 80 insertions(+) diff --git a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h index 7507453b00..762a0dc9a5 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h @@ -46,6 +46,7 @@ struct MnnvlSubBatch : public Transport::SubBatch { size_t max_size; CUDAStreamHandle sync_stream; CUDAStreamHandle async_stream; + int stream_device_id = -1; // Completion events created in startTransfer (one per submit). Destroyed by // the destructor (RAII); Slab::deallocate() invokes ~MnnvlSubBatch() // before reusing the storage, so this runs on every free. diff --git a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h index c619faec6c..fd786f08b2 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h @@ -48,6 +48,7 @@ struct NVLinkSubBatch : public Transport::SubBatch { size_t max_size; CUDAStreamHandle sync_stream; CUDAStreamHandle async_stream; + int stream_device_id = -1; // Completion events created in startTransfer (one per submit). Destroyed by // the destructor (RAII); Slab::deallocate() invokes ~NVLinkSubBatch() // before reusing the storage, so this runs on every free. diff --git a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp index ca5b8702b6..2e869018ea 100644 --- a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp @@ -238,6 +238,7 @@ Status MnnvlTransport::submitTransferTasks( stream_device)); CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->async_stream, stream_device)); + mnnvl_batch->stream_device_id = stream_device; } startTransfer(new_tasks, mnnvl_batch); @@ -312,11 +313,39 @@ void MnnvlTransport::startTransfer(std::vector &tasks, return; } + // Save and set device to match the stream's device to ensure event + // creation and recording happen on the correct device (fix for #2722). + int saved_device = -1; + if (batch->stream_device_id >= 0) { + auto err = cudaGetDevice(&saved_device); + if (err != cudaSuccess) { + LOG(ERROR) << "MnnvlTransport: cudaGetDevice failed: " + << cudaGetErrorString(err); + for (auto *task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + if (saved_device != batch->stream_device_id) { + err = cudaSetDevice(batch->stream_device_id); + if (err != cudaSuccess) { + LOG(ERROR) << "MnnvlTransport: cudaSetDevice failed: " + << cudaGetErrorString(err); + for (auto *task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + } + } + cudaEvent_t event; auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); if (event_err != cudaSuccess) { LOG(ERROR) << "MnnvlTransport: cudaEventCreateWithFlags failed: " << cudaGetErrorString(event_err); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } @@ -325,9 +354,19 @@ void MnnvlTransport::startTransfer(std::vector &tasks, LOG(ERROR) << "MnnvlTransport: cudaEventRecord failed: " << cudaGetErrorString(record_err); cudaEventDestroy(event); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } + + // Restore original device + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } + batch->completion_events.push_back(event); for (auto *task : tasks) task->completion_event = event; } diff --git a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp index 157c97f073..f3de90d7f3 100644 --- a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp @@ -194,6 +194,7 @@ Status NVLinkTransport::submitTransferTasks( stream_device)); CHECK_STATUS(platform_->getStreamFromPool(shm_batch->async_stream, stream_device)); + shm_batch->stream_device_id = stream_device; } startTransfer(new_tasks, shm_batch); @@ -273,11 +274,39 @@ void NVLinkTransport::startTransfer(std::vector& tasks, } } + // Save and set device to match the stream's device to ensure event + // creation and recording happen on the correct device (fix for #2722). + int saved_device = -1; + if (batch->stream_device_id >= 0) { + auto err = cudaGetDevice(&saved_device); + if (err != cudaSuccess) { + LOG(ERROR) << "NVLinkTransport: cudaGetDevice failed: " + << cudaGetErrorString(err); + for (auto* task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + if (saved_device != batch->stream_device_id) { + err = cudaSetDevice(batch->stream_device_id); + if (err != cudaSuccess) { + LOG(ERROR) << "NVLinkTransport: cudaSetDevice failed: " + << cudaGetErrorString(err); + for (auto* task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + } + } + cudaEvent_t event; auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); if (event_err != cudaSuccess) { LOG(ERROR) << "NVLinkTransport: cudaEventCreateWithFlags failed: " << cudaGetErrorString(event_err); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } @@ -286,9 +315,19 @@ void NVLinkTransport::startTransfer(std::vector& tasks, LOG(ERROR) << "NVLinkTransport: cudaEventRecord failed: " << cudaGetErrorString(record_err); cudaEventDestroy(event); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } + + // Restore original device + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } + batch->completion_events.push_back(event); for (auto* task : tasks) task->completion_event = event; } From af29ca7e02ef79342c8b2ee30c5b5c92136dc5ea Mon Sep 17 00:00:00 2001 From: Zhaoyi Li <36555117+Lzy17@users.noreply.github.com> Date: Mon, 6 Jul 2026 03:35:35 -0500 Subject: [PATCH 17/57] [TE] Route cross-host targets over rdma automatically in rdma+hip multi-protocol segments (#2753) * [TE] Route cross-host targets over rdma automatically in rdma+hip segments A "rdma,hip" multi-protocol segment registers the device KV pool under both rdma and hip. hip transport uses GPU IPC, which only works between processes on the same physical host, so selecting hip for a target on another host makes the initiator call hipIpcOpenMemHandle on a remote GPU handle and fail with "Error code 17 - invalid device pointer". Until now the only mitigation was the MC_DISABLE_HIP env, a global switch that also disables the intra-node hip fast path. Add an automatic same-host locality gate instead: compare the host portion of the target segment name against the local server name and skip hip buffers for cross-host targets so they fall back to rdma. Intra-node targets still use hip. mp_selectTransport mirrors the same downgrade for an explicit hip preference. The locality helpers are factored into multi_transport_locality.h with a unit test (multi_transport_locality_test), matching the rdma_gid_probe pattern. MC_DISABLE_HIP remains as a manual override for backward compatibility. * Make hip locality gate IPv6-safe and case-insensitive segmentHost now parses bracketed [v6]:port, bare IPv6 literals (multi-colon, no port -> whole string is the host), and the plain host:port / hostname forms. Host comparison is case-insensitive (hostEquals), since DNS names and IPv6 hex literals may differ only in letter case. mp_selectTransport now downgrades a cross-host hip preference to rdma, then tcp, and returns NotSupportedTransport when neither is offered, instead of silently keeping hip when rdma is absent. Adds SegmentHostParsesIPv6, HipReachableForIPv6, and HostMatchIsCaseInsensitive tests. --- .../include/multi_transport_locality.h | 82 ++++++++++++++++++ .../src/multi_transport.cpp | 38 +++++++++ mooncake-transfer-engine/tests/CMakeLists.txt | 7 ++ .../tests/multi_transport_locality_test.cpp | 84 +++++++++++++++++++ 4 files changed, 211 insertions(+) create mode 100644 mooncake-transfer-engine/include/multi_transport_locality.h create mode 100644 mooncake-transfer-engine/tests/multi_transport_locality_test.cpp diff --git a/mooncake-transfer-engine/include/multi_transport_locality.h b/mooncake-transfer-engine/include/multi_transport_locality.h new file mode 100644 index 0000000000..c298b4a702 --- /dev/null +++ b/mooncake-transfer-engine/include/multi_transport_locality.h @@ -0,0 +1,82 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MULTI_TRANSPORT_LOCALITY_H +#define MULTI_TRANSPORT_LOCALITY_H + +#include +#include + +namespace mooncake { + +// Extract the host portion of a segment name. Segment names carry an optional +// ":port" suffix, and the host may be an IPv4 address, a hostname, or an IPv6 +// literal. IPv6 literals contain multiple colons, so a naive rfind(':') would +// corrupt them; the following forms are handled explicitly: +// "10.0.0.1:8000" -> "10.0.0.1" (IPv4 / hostname with port) +// "node-a" -> "node-a" (no port) +// "[2001:db8::1]:8000" -> "2001:db8::1" (bracketed IPv6 with port) +// "[2001:db8::1]" -> "2001:db8::1" (bracketed IPv6, no port) +// "2001:db8::1" -> "2001:db8::1" (bare IPv6, no port) +inline std::string segmentHost(const std::string& segment_name) { + if (!segment_name.empty() && segment_name.front() == '[') { + // Bracketed IPv6 literal: strip the brackets and ignore any ":port". + auto close = segment_name.find(']'); + if (close != std::string::npos) { + return segment_name.substr(1, close - 1); + } + return segment_name; // Malformed; return as-is. + } + auto first = segment_name.find(':'); + if (first == std::string::npos) { + return segment_name; // No port and no colon. + } + if (first != segment_name.rfind(':')) { + // More than one colon and not bracketed: a bare IPv6 literal without a + // port (an IPv6 host with a port must use brackets). Treat it all as + // the host. + return segment_name; + } + return segment_name.substr(0, first); // "host:port". +} + +// Case-insensitive equality. Hostnames are case-insensitive per DNS, and IPv6 +// hex literals may differ only in letter case, so a plain "==" would spuriously +// classify the same host as remote and lose the intra-node hip fast path. +inline bool hostEquals(const std::string& a, const std::string& b) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) { + if (std::tolower(static_cast(a[i])) != + std::tolower(static_cast(b[i]))) { + return false; + } + } + return true; +} + +// The device KV pool is registered under both rdma and hip in a "rdma,hip" +// multi-protocol segment. hip transport uses GPU IPC, which only works between +// processes on the same physical host, so a hip buffer is a valid transport for +// a target only when that target lives on the same host as the initiator. +// A cross-host target must fall back to rdma. Two engines co-located on one +// host share the host portion of their segment name (they differ only in port). +inline bool isHipReachableTarget(const std::string& target_segment_name, + const std::string& local_server_name) { + return hostEquals(segmentHost(target_segment_name), + segmentHost(local_server_name)); +} + +} // namespace mooncake + +#endif // MULTI_TRANSPORT_LOCALITY_H diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 30475ede7d..ecf18e5ebd 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -19,6 +19,7 @@ #include #include "config.h" +#include "multi_transport_locality.h" #include "transport/rdma_transport/rdma_transport.h" #ifdef USE_BAREX #include "transport/barex_transport/barex_transport.h" @@ -475,6 +476,14 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, if (p == "tcp") return 1; return 0; }; + // hip transport uses GPU IPC, which cannot reach a GPU on another host. + // The device KV pool is registered under both rdma and hip, so a + // cross-host target must skip its hip buffers and fall back to rdma. + // This makes the intra-node fast path (hip) and the cross-node path + // (rdma) work automatically from a single multi-protocol segment, + // without requiring the operator to set MC_DISABLE_HIP. + const bool hip_reachable = + isHipReachableTarget(target_segment_desc->name, local_server_name_); std::string chosen; int chosen_priority = -1; for (const auto& buffer : target_segment_desc->buffers) { @@ -486,6 +495,7 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, : buffer.addr; if (entry.target_offset >= start && entry.target_offset < start + buffer.length) { + if (buffer.protocol == "hip" && !hip_reachable) continue; int priority = protocol_priority(buffer.protocol); if (priority > chosen_priority) { chosen = buffer.protocol; @@ -503,6 +513,12 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, return Status::NotSupportedTransport("Transport " + chosen + " not installed"); } + if (globalConfig().trace) { + LOG(INFO) << "MultiTransport::selectTransport route: target_id=" + << entry.target_id << " segment_protocol=\"" << proto + << "\" hip_reachable=" << hip_reachable + << " chosen=" << chosen; + } transport = transport_map_[chosen].get(); return Status::OK(); } @@ -540,6 +556,28 @@ Status MultiTransport::mp_selectTransport(const TransferRequest& entry, if (!item.empty()) protos.push_back(item); } + // hip GPU IPC cannot reach a remote host; downgrade an explicit hip + // preference to a cross-host-capable transport for a cross-host target + // (mirrors the locality gate in selectTransport). Prefer rdma, then tcp. + if (preferred_proto == "hip" && + !isHipReachableTarget(target_segment_desc->name, local_server_name_)) { + std::string fallback; + for (const char* candidate : {"rdma", "tcp"}) { + if (std::find(protos.begin(), protos.end(), candidate) != + protos.end()) { + fallback = candidate; + break; + } + } + if (fallback.empty()) { + return Status::NotSupportedTransport( + "hip target is cross-host but segment " + + std::to_string(entry.target_id) + + " offers no cross-host transport (rdma/tcp)"); + } + preferred_proto = fallback; + } + #ifdef USE_ASCEND_HETEROGENEOUS // When USE_ASCEND_HETEROGENEOUS is enabled: // - Target side directly reuses RDMA Transport diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 1cca40a7da..83b35b353b 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -244,6 +244,13 @@ target_link_libraries(rdma_gid_probe_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME rdma_gid_probe_test COMMAND rdma_gid_probe_test) +add_executable(multi_transport_locality_test + ${WORKSPACE}/multi_transport_locality_test.cpp) +target_link_libraries(multi_transport_locality_test PUBLIC transfer_engine gtest + gtest_main) +add_test(NAME multi_transport_locality_test + COMMAND multi_transport_locality_test) + add_executable(rdma_context_reprobe_test ${WORKSPACE}/rdma_context_reprobe_test.cpp) target_link_libraries(rdma_context_reprobe_test PUBLIC transfer_engine gtest diff --git a/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp b/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp new file mode 100644 index 0000000000..750d9e60bd --- /dev/null +++ b/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp @@ -0,0 +1,84 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include + +#include "multi_transport_locality.h" + +using namespace mooncake; + +TEST(MultiTransportLocalityTest, SegmentHostStripsPort) { + EXPECT_EQ(segmentHost("10.0.0.1:12345"), "10.0.0.1"); + EXPECT_EQ(segmentHost("node-a:8000"), "node-a"); + // No port: whole string is the host. + EXPECT_EQ(segmentHost("node-a"), "node-a"); + EXPECT_EQ(segmentHost(""), ""); +} + +TEST(MultiTransportLocalityTest, HipReachableForSameHostDifferentPort) { + // Two engines co-located on one host share the host, differ only in port. + EXPECT_TRUE(isHipReachableTarget("10.0.0.1:20000", "10.0.0.1:20001")); + EXPECT_TRUE(isHipReachableTarget("node-a:8000", "node-a:9000")); +} + +TEST(MultiTransportLocalityTest, HipReachableForIdenticalName) { + EXPECT_TRUE(isHipReachableTarget("node-a:8000", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, HipNotReachableForRemoteHost) { + // Cross-host target: hip (GPU IPC) is not usable, must fall back to rdma. + EXPECT_FALSE(isHipReachableTarget("10.0.0.2:20000", "10.0.0.1:20000")); + EXPECT_FALSE(isHipReachableTarget("node-b:8000", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, HandlesMissingPort) { + // Locality decision still works when one side has no explicit port. + EXPECT_TRUE(isHipReachableTarget("node-a", "node-a:8000")); + EXPECT_FALSE(isHipReachableTarget("node-b", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, SegmentHostParsesIPv6) { + // Bracketed IPv6 with and without a port. + EXPECT_EQ(segmentHost("[2001:db8::1]:8000"), "2001:db8::1"); + EXPECT_EQ(segmentHost("[2001:db8::1]"), "2001:db8::1"); + EXPECT_EQ(segmentHost("[::1]:20000"), "::1"); + // Bare IPv6 literal without a port: the whole string is the host. + EXPECT_EQ(segmentHost("2001:db8::1"), "2001:db8::1"); + EXPECT_EQ(segmentHost("::1"), "::1"); +} + +TEST(MultiTransportLocalityTest, HipReachableForIPv6) { + // Same IPv6 host, different ports (bracketed) -> intra-node hip. + EXPECT_TRUE( + isHipReachableTarget("[2001:db8::1]:20000", "[2001:db8::1]:20001")); + // Bracketed-with-port vs bare literal for the same host must still match. + EXPECT_TRUE(isHipReachableTarget("[2001:db8::1]:20000", "2001:db8::1")); + // Different IPv6 hosts -> cross-node, fall back to rdma. + EXPECT_FALSE( + isHipReachableTarget("[2001:db8::1]:20000", "[2001:db8::2]:20000")); +} + +TEST(MultiTransportLocalityTest, HostMatchIsCaseInsensitive) { + // Hostnames and IPv6 hex literals are case-insensitive. + EXPECT_TRUE(isHipReachableTarget("Node-A:8000", "node-a:9000")); + EXPECT_TRUE( + isHipReachableTarget("[2001:DB8::1]:8000", "[2001:db8::1]:9000")); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} From d4cfe3c8c0702a7e2f0beee888cc12b18da428a7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:54:16 +0800 Subject: [PATCH 18/57] Bump golang.org/x/net (#2730) Bumps [golang.org/x/net](https://github.com/golang/net) from 0.38.0 to 0.55.0. - [Commits](https://github.com/golang/net/compare/v0.38.0...v0.55.0) --- updated-dependencies: - dependency-name: golang.org/x/net dependency-version: 0.55.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../example/http-metadata-server/go.mod | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/mooncake-transfer-engine/example/http-metadata-server/go.mod b/mooncake-transfer-engine/example/http-metadata-server/go.mod index 84a942aa88..2646bb118a 100644 --- a/mooncake-transfer-engine/example/http-metadata-server/go.mod +++ b/mooncake-transfer-engine/example/http-metadata-server/go.mod @@ -1,8 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-transfer-engine/example/http-metadata-server -go 1.23.0 - -toolchain go1.23.8 +go 1.25.0 require github.com/gin-gonic/gin v1.10.0 @@ -27,10 +25,10 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect golang.org/x/arch v0.8.0 // indirect - golang.org/x/crypto v0.36.0 // indirect - golang.org/x/net v0.38.0 // indirect - golang.org/x/sys v0.31.0 // indirect - golang.org/x/text v0.23.0 // indirect + golang.org/x/crypto v0.51.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect google.golang.org/protobuf v1.34.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) From f1f5e40b82159247fa10f244724c6e8d8d7b98e0 Mon Sep 17 00:00:00 2001 From: Stary Date: Mon, 6 Jul 2026 17:26:21 +0800 Subject: [PATCH 19/57] [CI/Build] Add dedicated tent-ci job with CUDA/CPU matrix (#2720) * [CI/Build] Add dedicated tent-ci job with CUDA/CPU matrix Extract TENT build/test into a dedicated job with a USE_CUDA=ON / USE_CUDA=OFF matrix, gate it on transfer-engine/common/CMake path changes, and require it in ci-gate. Signed-off-by: staryxchen * [CI/Build] Skip TENT tests on cuda-on leg (build-only) GitHub runners have no real GPU; with USE_CUDA=ON tent's cuda_probe hits the CUDA stub driver at runtime and drives dispatch past the fake objects the unit tests rely on, causing 6 tests in 4 binaries to fail. Restrict Test (TENT) to the cuda-off leg. cuda-on still compiles every #ifdef USE_CUDA branch. Signed-off-by: staryxchen --------- Signed-off-by: staryxchen --- .github/workflows/ci.yml | 133 ++++++++++++++++++++++++++++++++------- 1 file changed, 112 insertions(+), 21 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 431dbd5895..3c3cad1dea 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -729,26 +729,6 @@ jobs: sudo cmake --install . shell: bash - - name: Configure project with TENT - run: | - mkdir build-tent - cd build-tent - cmake -G Ninja .. -DUSE_TENT=ON -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_DEBUG_SYMBOLS=OFF - shell: bash - - - name: Build project with TENT - run: | - cd build-tent - cmake --build . - sudo cmake --install . - shell: bash - - - name: Test (TENT) - run: | - cd build-tent - ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure - shell: bash - - name: Build nvlink_allocator.so run: | mkdir -p build/mooncake-transfer-engine/nvlink-allocator @@ -925,12 +905,15 @@ jobs: runs-on: ubuntu-latest outputs: should-run-downstream: ${{ steps.dispatch-override.outputs.src || steps.filter.outputs.src }} + should-run-tent: ${{ steps.dispatch-override.outputs.tent || steps.filter.outputs.tent }} steps: # workflow_dispatch has no PR/push diff context — skip paths-filter and default to true - name: Default to true for workflow_dispatch id: dispatch-override if: github.event_name == 'workflow_dispatch' - run: echo "src=true" >> $GITHUB_OUTPUT + run: | + echo "src=true" >> $GITHUB_OUTPUT + echo "tent=true" >> $GITHUB_OUTPUT - uses: actions/checkout@v4 if: github.event_name != 'workflow_dispatch' with: @@ -948,6 +931,12 @@ jobs: - 'dependencies.sh' - 'scripts/**' - '.github/workflows/**' + tent: + - 'mooncake-transfer-engine/**' + - 'mooncake-common/**' + - 'CMakeLists.txt' + - 'dependencies.sh' + - '.github/workflows/ci.yml' build-wheel-cu13: needs: [spell-check, clang-format, check-paths] @@ -985,6 +974,107 @@ jobs: uses: ./.github/workflows/integration-test.yml secrets: inherit + tent-ci: + needs: [spell-check, clang-format, check-paths] + if: >- + (needs.check-paths.outputs.should-run-tent == 'true' || + github.event_name == 'workflow_dispatch') && + (github.event_name == 'push' || + github.event_name == 'workflow_dispatch' || + github.event.action == 'opened' || + contains(github.event.pull_request.labels.*.name, 'run-ci')) + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + include: + - name: cuda-on + cmake_flags: '-DUSE_CUDA=ON -DCMAKE_EXE_LINKER_FLAGS=-L/usr/local/cuda/lib64/stubs' + need_cuda: true + - name: cuda-off + cmake_flags: '-DUSE_CUDA=OFF' + need_cuda: false + name: tent-ci (${{ matrix.name }}) + env: + CI: "true" + SCCACHE_GHA_ENABLED: "true" + + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Free up disk space + if: matrix.need_cuda + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf /opt/hostedtoolcache/CodeQL + sudo rm -rf /usr/local/lib/android + df -h + + - name: Install CUDA Toolkit + if: matrix.need_cuda + uses: Jimver/cuda-toolkit@v0.2.24 + with: + cuda: '12.8.1' + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc", "nvrtc-dev"]' + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Install dependencies + run: | + sudo apt update -y + sudo apt install -y ninja-build + sudo bash -x dependencies.sh -y + df -h + shell: bash + + - name: Configure project with TENT + run: | + mkdir build-tent + cd build-tent + cmake -G Ninja .. -DUSE_TENT=ON -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_DEBUG_SYMBOLS=OFF ${{ matrix.cmake_flags }} + shell: bash + + - name: Build project with TENT + run: | + if [ "${{ matrix.need_cuda }}" = "true" ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH + fi + cd build-tent + cmake --build . + sudo cmake --install . + shell: bash + + # Only run tests on the cuda-off leg. GitHub runners have no real GPU; + # with USE_CUDA=ON tent's cuda_probe hits the CUDA stub library at + # runtime and drives some dispatch paths past the fake objects the + # unit tests rely on, causing false failures. cuda-on still validates + # that every #ifdef USE_CUDA branch compiles. + - name: Test (TENT) + if: '!matrix.need_cuda' + run: | + cd build-tent + ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure + shell: bash + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + shell: bash + run: ${SCCACHE_PATH} --show-stats + ci-gate: name: CI Gate if: always() @@ -1000,6 +1090,7 @@ jobs: - test-wheel-ubuntu - build-wheel-cu13 - build-wheel-efa + - tent-ci - ascend-test - integration-test runs-on: ubuntu-latest From a407c52581f12b92a04e03897f4cf6efbc2ca565 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Tue, 7 Jul 2026 01:23:51 +0800 Subject: [PATCH 20/57] [Doc] Update README with LightX2V deployment details (#2767) Added details about LightX2V's support for disaggregated deployment and linked to the relevant blog post. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b08dbfee34..107e782c05 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ Under real workloads, Mooncake’s innovative architecture enables Kimi to handl - **May 7, 2026**: 🚀 [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) — a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing! - **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s → 7.2s) with zero-copy RDMA transfer across thousands of GPUs. - **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management. -- **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer. +- **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer. Details in [blog](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/). - **Feb 25, 2026**: [SGLang](https://github.com/sgl-project/sglang) merged [Encoder Global Cache Manager](https://github.com/sgl-project/sglang/pull/16137), introducing a Mooncake-powered global multimodal embedding cache that enables cross-instance sharing of ViT embeddings to avoid redundant GPU computation.
From 7c03cc9a4b557e0d449bfad2f504f4113aa80693 Mon Sep 17 00:00:00 2001 From: Posedge_Lin Date: Tue, 7 Jul 2026 10:24:40 +0800 Subject: [PATCH 21/57] [Store] Add BatchEvict candidate-selection benchmark (#2584) --- .github/workflows/ci.yml | 2 +- mooncake-store/benchmarks/CMakeLists.txt | 6 + .../benchmarks/batch_evict_bench.cpp | 437 ++++++++++++++++++ mooncake-store/include/master_service.h | 4 + 4 files changed, 448 insertions(+), 1 deletion(-) create mode 100644 mooncake-store/benchmarks/batch_evict_bench.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3c3cad1dea..bb2e9925b7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -228,7 +228,7 @@ jobs: } echo "=== Processing coverage data ===" - lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' --output-file coverage.filtered.info 2>&1 || true + lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' '*/benchmarks/*' --output-file coverage.filtered.info 2>&1 || true echo "=== Generating HTML report ===" genhtml coverage.filtered.info --output-directory coverage_report 2>&1 || echo "genhtml failed, continuing..." diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt index 76d23cda09..98769bd937 100644 --- a/mooncake-store/benchmarks/CMakeLists.txt +++ b/mooncake-store/benchmarks/CMakeLists.txt @@ -39,3 +39,9 @@ if(USE_NOF) target_link_libraries(nof_worker_pool_bench PRIVATE mooncake_store glog::glog gflags::gflags) endif() + +# Add BatchEvict benchmark executable +add_executable(batch_evict_bench batch_evict_bench.cpp) +target_link_libraries( + batch_evict_bench PRIVATE mooncake_store cachelib_memory_allocator + gflags::gflags glog::glog pthread) diff --git a/mooncake-store/benchmarks/batch_evict_bench.cpp b/mooncake-store/benchmarks/batch_evict_bench.cpp new file mode 100644 index 0000000000..51949e8e94 --- /dev/null +++ b/mooncake-store/benchmarks/batch_evict_bench.cpp @@ -0,0 +1,437 @@ +#include "master_service.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" +#include "types.h" + +namespace mooncake::benchmarks { + +class BatchEvictBench { + public: + static bool RunRealBatchEvictScales() { + std::vector scales = {10000, 100000}; + const char* large_mode = std::getenv("MOONCAKE_EVICT_BENCH_LARGE"); + if (large_mode != nullptr && std::string(large_mode) == "1") { + scales.push_back(1000000); + } + + std::cout << "num_objects,total_us,objects_before,objects_after," + "evicted_count,freed_bytes" + << std::endl; + for (size_t scale : scales) { + if (!RunOneScale(scale)) { + return false; + } + } + return true; + } + + static bool RunSingleWaiterSnapshotMutexProbe() { + if (std::getenv("MOONCAKE_EVICT_BENCH_LOCK_PROBE") == nullptr) { + return true; + } + + const size_t num_objects = + ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_OBJECTS", 1000000); + const size_t trials = std::max( + 1, ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_TRIALS", 30)); + const auto waiter_delay = std::chrono::microseconds( + static_cast( + ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_DELAY_US", 1000))); + + std::vector batch_evict_total_us; + std::vector unique_lock_wait_us; + + batch_evict_total_us.reserve(trials); + unique_lock_wait_us.reserve(trials); + + for (size_t trial = 0; trial < trials; ++trial) { + LockProbeTrialResult result; + if (!RunLockProbeTrial(num_objects, waiter_delay, result)) { + return false; + } + + batch_evict_total_us.push_back(result.batch_evict_total_us); + unique_lock_wait_us.push_back(result.unique_lock_wait_us); + } + + std::cout << "num_objects,trials,batch_evict_total_p50_us," + "unique_lock_wait_p50_us,unique_lock_wait_p95_us," + "unique_lock_wait_max_us" + << std::endl; + + std::cout << num_objects << "," << trials << "," + << PercentileValue(batch_evict_total_us, 0.50) << "," + << PercentileValue(unique_lock_wait_us, 0.50) << "," + << PercentileValue(unique_lock_wait_us, 0.95) << "," + << PercentileValue(unique_lock_wait_us, 1.00) << std::endl; + return true; + } + + private: + static constexpr const char* kTenantId = "default"; + static constexpr const char* kSegmentName = "batch_evict_bench_segment"; + static constexpr size_t kSegmentBase = 0x300000000; + static constexpr uint64_t kObjectSize = 1024; + static constexpr double kEvictRatioTarget = 0.50; + static constexpr double kEvictRatioLowerbound = 0.25; + + struct MetadataStats { + size_t object_count{0}; + size_t completed_memory_replicas{0}; + size_t busy_memory_replicas{0}; + size_t non_memory_replicas{0}; + size_t incomplete_replicas{0}; + size_t unexpired_leases{0}; + }; + + struct LockProbeTrialResult { + uint64_t batch_evict_total_us{0}; + uint64_t unique_lock_wait_us{0}; + }; + + static MasterServiceConfig MakeConfig() { + return MasterServiceConfig::builder() + .set_memory_allocator(BufferAllocatorType::OFFSET) + .set_eviction_ratio(0.0) + .set_eviction_high_watermark_ratio(1.0) + .set_client_live_ttl_sec(3600) + .build(); + } + + static size_t SegmentSizeFor(size_t num_objects) { + constexpr size_t kMinSegmentSize = 16 * 1024 * 1024; + const size_t needed = num_objects * kObjectSize; + const size_t headroom = needed / 8 + 1024 * kObjectSize; + return std::max(kMinSegmentSize, needed + headroom); + } + + static Segment MakeSegment(size_t num_objects) { + Segment segment; + segment.id = generate_uuid(); + segment.name = kSegmentName; + segment.base = kSegmentBase; + segment.size = SegmentSizeFor(num_objects); + segment.te_endpoint = segment.name; + return segment; + } + + static std::string MakeKey(size_t index) { + return "batch_evict_bench_key_" + std::to_string(index); + } + + static size_t ReadEnvSize(const char* name, size_t default_value) { + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return default_value; + } + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + if (end == value) { + return default_value; + } + return static_cast(parsed); + } + + static uint64_t PercentileValue(std::vector values, + double percentile) { + if (values.empty()) { + return 0; + } + std::sort(values.begin(), values.end()); + const size_t rank = std::max( + 1, static_cast(std::ceil(percentile * values.size()))); + return values[std::min(rank - 1, values.size() - 1)]; + } + + static MetadataStats ExpireLeasesAndCollectStats(MasterService& service) { + MetadataStats stats; + auto now = std::chrono::system_clock::now(); + const auto base_expiration = now - std::chrono::hours(1); + size_t ordinal = 0; + + for (size_t shard_idx = 0; shard_idx < MasterService::kNumShards; + ++shard_idx) { + MasterService::MetadataShardAccessorRW shard(&service, shard_idx); + for (auto& [tenant_id, tenant_state] : shard->tenants) { + if (tenant_id != kTenantId) { + continue; + } + for (auto& [key, metadata] : tenant_state.metadata) { + { + SpinLocker locker(&metadata.lock); + metadata.lease_timeout = + base_expiration + + std::chrono::nanoseconds(ordinal++); + } + + ++stats.object_count; + if (!metadata.IsLeaseExpired(now)) { + ++stats.unexpired_leases; + } + for (const auto& replica : metadata.GetAllReplicas()) { + if (replica.is_memory_replica()) { + if (replica.is_completed()) { + ++stats.completed_memory_replicas; + } else { + ++stats.incomplete_replicas; + } + if (replica.get_refcnt() != 0) { + ++stats.busy_memory_replicas; + } + } else { + ++stats.non_memory_replicas; + } + } + } + } + } + + return stats; + } + + static bool MountBenchSegment(MasterService& service, const UUID& client_id, + size_t num_objects) { + auto segment = MakeSegment(num_objects); + auto mount_result = service.MountSegment(segment, client_id); + if (!mount_result.has_value()) { + LOG(ERROR) << "MountSegment failed: " + << toString(mount_result.error()); + return false; + } + return true; + } + + static bool CreateCompletedMemoryObjects(MasterService& service, + const UUID& client_id, + size_t num_objects) { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = kSegmentName; + + for (size_t i = 0; i < num_objects; ++i) { + const std::string key = MakeKey(i); + auto put_start = service.PutStart(client_id, key, kTenantId, + kObjectSize, config); + if (!put_start.has_value()) { + LOG(ERROR) << "PutStart failed for i=" << i + << ", error=" << toString(put_start.error()); + return false; + } + if (put_start->size() != 1u) { + LOG(ERROR) << "PutStart returned " << put_start->size() + << " replicas for i=" << i; + return false; + } + + auto put_end = + service.PutEnd(client_id, key, kTenantId, ReplicaType::MEMORY); + if (!put_end.has_value()) { + LOG(ERROR) << "PutEnd failed for i=" << i + << ", error=" << toString(put_end.error()); + return false; + } + } + return true; + } + + static bool UsedBytes(MasterService& service, uint64_t& used_bytes) { + auto segment_usage = service.QuerySegments(kSegmentName); + if (!segment_usage.has_value()) { + LOG(ERROR) << "QuerySegments failed: " + << toString(segment_usage.error()); + return false; + } + used_bytes = segment_usage->first; + return true; + } + + static uint64_t WaitForSnapshotUniqueLock( + MasterService& service, std::chrono::microseconds waiter_delay) { + std::this_thread::sleep_for(waiter_delay); + const auto wait_start = std::chrono::steady_clock::now(); + uint64_t wait_us = 0; + { + std::unique_lock lock(service.snapshot_mutex_); + const auto wait_end = std::chrono::steady_clock::now(); + wait_us = std::chrono::duration_cast( + wait_end - wait_start) + .count(); + } + return wait_us; + } + + static bool ValidatePopulatedObjects(const MasterService& service, + size_t num_objects, + const MetadataStats& stats) { + if (service.GetKeyCount() != num_objects) { + LOG(ERROR) << "GetKeyCount mismatch: expected=" << num_objects + << ", actual=" << service.GetKeyCount(); + return false; + } + if (stats.object_count != num_objects || + stats.completed_memory_replicas != num_objects || + stats.busy_memory_replicas != 0 || stats.non_memory_replicas != 0 || + stats.incomplete_replicas != 0 || stats.unexpired_leases != 0) { + LOG(ERROR) << "metadata validation failed: objects=" + << stats.object_count << ", completed_memory_replicas=" + << stats.completed_memory_replicas + << ", busy_memory_replicas=" + << stats.busy_memory_replicas + << ", non_memory_replicas=" << stats.non_memory_replicas + << ", incomplete_replicas=" << stats.incomplete_replicas + << ", unexpired_leases=" << stats.unexpired_leases; + return false; + } + return true; + } + + static bool ValidateEvictionResult(size_t objects_before, + size_t evicted_count, + uint64_t freed_bytes) { + const size_t lowerbound = static_cast( + std::ceil(objects_before * kEvictRatioLowerbound)); + if (evicted_count < lowerbound) { + LOG(ERROR) << "evicted_count below lowerbound: evicted=" + << evicted_count << ", lowerbound=" << lowerbound; + return false; + } + if (evicted_count * kObjectSize != freed_bytes) { + LOG(ERROR) << "freed_bytes mismatch: evicted_count=" + << evicted_count << ", object_size=" << kObjectSize + << ", freed_bytes=" << freed_bytes; + return false; + } + return true; + } + + static bool RunOneScale(size_t num_objects) { + MasterService service(MakeConfig()); + const UUID client_id = generate_uuid(); + + if (!MountBenchSegment(service, client_id, num_objects) || + !CreateCompletedMemoryObjects(service, client_id, num_objects)) { + return false; + } + + const MetadataStats stats = ExpireLeasesAndCollectStats(service); + if (!ValidatePopulatedObjects(service, num_objects, stats)) { + return false; + } + + const size_t objects_before = service.GetKeyCount(); + uint64_t used_before = 0; + if (!UsedBytes(service, used_before)) { + return false; + } + + const auto evict_start = std::chrono::steady_clock::now(); + service.BatchEvict(kEvictRatioTarget, kEvictRatioLowerbound); + const auto total_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - evict_start) + .count(); + + const size_t objects_after = service.GetKeyCount(); + uint64_t used_after = 0; + if (!UsedBytes(service, used_after)) { + return false; + } + const size_t evicted_count = objects_before - objects_after; + const uint64_t freed_bytes = + used_before >= used_after ? used_before - used_after : 0; + + if (!ValidateEvictionResult(objects_before, evicted_count, + freed_bytes)) { + return false; + } + + std::cout << num_objects << "," << total_us << "," << objects_before + << "," << objects_after << "," << evicted_count << "," + << freed_bytes << std::endl; + return true; + } + + static bool RunLockProbeTrial(size_t num_objects, + std::chrono::microseconds waiter_delay, + LockProbeTrialResult& result) { + MasterService service(MakeConfig()); + const UUID client_id = generate_uuid(); + + if (!MountBenchSegment(service, client_id, num_objects) || + !CreateCompletedMemoryObjects(service, client_id, num_objects)) { + return false; + } + + const MetadataStats stats = ExpireLeasesAndCollectStats(service); + if (!ValidatePopulatedObjects(service, num_objects, stats)) { + return false; + } + + const size_t objects_before = service.GetKeyCount(); + uint64_t used_before = 0; + if (!UsedBytes(service, used_before)) { + return false; + } + + uint64_t unique_lock_wait_us = 0; + std::thread waiter([&]() { + unique_lock_wait_us = + WaitForSnapshotUniqueLock(service, waiter_delay); + }); + + const auto evict_start = std::chrono::steady_clock::now(); + service.BatchEvict(kEvictRatioTarget, kEvictRatioLowerbound); + const auto batch_evict_total_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - evict_start) + .count(); + waiter.join(); + + const size_t objects_after = service.GetKeyCount(); + uint64_t used_after = 0; + if (!UsedBytes(service, used_after)) { + return false; + } + const size_t evicted_count = objects_before - objects_after; + const uint64_t freed_bytes = + used_before >= used_after ? used_before - used_after : 0; + + if (!ValidateEvictionResult(objects_before, evicted_count, + freed_bytes)) { + return false; + } + + result.batch_evict_total_us = + static_cast(batch_evict_total_us); + result.unique_lock_wait_us = unique_lock_wait_us; + return true; + } +}; + +} // namespace mooncake::benchmarks + +int main(int argc, char** argv) { + google::InitGoogleLogging("BatchEvictBench"); + FLAGS_logtostderr = true; + gflags::ParseCommandLineFlags(&argc, &argv, true); + + using mooncake::benchmarks::BatchEvictBench; + const bool ok = BatchEvictBench::RunRealBatchEvictScales() && + BatchEvictBench::RunSingleWaiterSnapshotMutexProbe(); + + google::ShutdownGoogleLogging(); + return ok ? 0 : 1; +} diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index b5591d6507..d7823d2468 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -62,6 +62,9 @@ class SnapshotChildProcessTest; class PromotionOnHitTest; class MasterServiceTenantQuotaTest; } // namespace test +namespace benchmarks { +class BatchEvictBench; +} // namespace benchmarks /* * @brief MasterService is the main class for the master server. @@ -82,6 +85,7 @@ class MasterService { friend class test::MasterServiceSnapshotTestBase; friend class test::SnapshotChildProcessTest; friend class test::PromotionOnHitTest; + friend class benchmarks::BatchEvictBench; friend class test::MasterServiceTenantQuotaTest; public: From dd4f40d6f2ff4014dc6a9138f0453ce5a2c5359a Mon Sep 17 00:00:00 2001 From: "Guocheng(Eric) Song" Date: Tue, 7 Jul 2026 10:33:34 +0800 Subject: [PATCH 22/57] docs: clarify local SSD offload configuration (#2728) --- .../source/deployment/mooncake-store-deployment-guide.md | 9 +++++++-- docs/source/deployment/ssd-offload.md | 4 +++- docs/source/design/mooncake-store.md | 2 ++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 9ad0f3aebd..9798cf3d23 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -199,11 +199,12 @@ mooncake_master \ --offload_on_evict=true \ --promotion_on_hit=true \ --promotion_admission_threshold=2 \ - --root_fs_dir=/mnt/ssd_cache \ --enable_http_metadata_server=true \ --http_metadata_server_port=8080 ``` +Do not set `--root_fs_dir` with `--enable_offload=true`. `--root_fs_dir` is a legacy parameter from an older persistence path and may cause issues on the SSD offload path. Configure each real client's offload directory with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` instead. + --- ### CXL-Aware Allocation — Memory Tiering @@ -554,6 +555,8 @@ Flags for controlling data movement between DRAM and SSD. Start with `--enable_offload=true` for eager asynchronous SSD persistence after `Put` completion. Add `--offload_on_evict=true` when you want SSD writes to happen only when memory pressure selects an object for eviction. Add `--promotion_on_hit=true` to allow hot SSD-only data to be promoted back to DRAM, and tune `--promotion_admission_threshold` to control how many observed reads are required before promotion is queued. +For SSD offload, configure the disk path on each real client with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`; the master tracks these objects as `LOCAL_DISK` replicas. Do not use the legacy `--root_fs_dir` parameter with `--enable_offload=true`. + When `--offload_on_evict=true` is active, each `BatchEvict` cycle can queue at most `offloading_queue_limit * offload_cap_ratio` objects for SSD offload (default: `50000 * 0.5 = 25000`); objects exceeding this cap fall back to force-evict (discard) if `--offload_force_evict=true`, otherwise they remain in memory. For SSD-heavy workloads where NVMe bandwidth is underutilized while the KV-cache hit rate suffers, raise both `--offloading_queue_limit` and `--offload_cap_ratio` so more objects per cycle are actually persisted to SSD instead of discarded. Example: `--offloading_queue_limit=500000 --offload_cap_ratio=0.8` yields a per-cycle cap of `400000` (vs the default `25000`). ### CXL Memory @@ -570,9 +573,11 @@ When `--allocation_strategy=cxl` is set alongside `--enable_cxl=true`, the maste | Flag | Default | Description | |------|---------|-------------| -| `--root_fs_dir` | empty | DFS mount directory for multi-layer storage backend | +| `--root_fs_dir` | empty | Legacy DFS persistence directory; do not use with SSD offload | | `--global_file_segment_size` | `INT64_MAX` (unlimited) | Max available space for DFS segments; default does not cap DFS usage | +`--root_fs_dir` is a legacy persistence parameter and is expected to be replaced as the distributed filesystem path is refactored. For SSD offload, configure `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` on each real client instead. + ### NoF (NVMe-oF SSD Pool) ```{caution} diff --git a/docs/source/deployment/ssd-offload.md b/docs/source/deployment/ssd-offload.md index 439f5f29b5..0a5aedd08d 100644 --- a/docs/source/deployment/ssd-offload.md +++ b/docs/source/deployment/ssd-offload.md @@ -2,7 +2,7 @@ ## Overview -Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from SSD. +Mooncake Store supports offloading KV cache objects from distributed memory to a local filesystem path, typically backed by local SSDs. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from the local filesystem-backed offload path. For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md). @@ -13,6 +13,8 @@ SSD offload requires the **Real Client** and supports two deployment modes: In both modes, all SSD reads and writes happen within the Real Client (embedded or standalone). +SSD offload does not use the master's `--root_fs_dir` option. Configure the local disk path on each Real Client with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`; the master tracks offloaded objects as `LOCAL_DISK` replicas. `--root_fs_dir` is a legacy parameter from an older persistence path and may cause issues when used with `--enable_offload=true`. + ## Startup Steps ### Step 1: Create the SSD storage directory diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md index 72912c20c9..6e02e3c8b9 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -731,6 +731,8 @@ When the user specifies `--root_fs_dir=/path/to/dir` when starting the master, a ​Note​​: When enabling this feature, the user must ensure that the DFS-mounted directory (`root_fs_dir=/path/to/dir`) is valid and consistent across all client hosts. If some clients have invalid or incorrect mount paths, it may cause abnormal behavior in Mooncake Store. +This `root_fs_dir` path is a legacy persistence path. SSD offload uses `--enable_offload=true` on the master and real client, stores data under the real client's `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`, and records `LOCAL_DISK` replicas. Do not use `--root_fs_dir` with `--enable_offload=true`. + ### Persistent Storage Space Configuration​ Mooncake provides configurable DFS available space. Users can specify `--global_file_segment_size=1048576` when starting the master, indicating a maximum usable space of 1MB on DFS. The current default setting is the maximum value of int64 (as we generally do not restrict DFS storage usage), which is displayed as `infinite` in `mooncake_maseter`'s console logs. From cf635de1e04f3d55a4d4b4247f0c28caee7b51af Mon Sep 17 00:00:00 2001 From: Leoy Date: Tue, 7 Jul 2026 10:34:38 +0800 Subject: [PATCH 23/57] [Wheel] Validate MooncakeConfig fields with fail-fast errors (#2456) --- mooncake-wheel/mooncake/mooncake_config.py | 111 +++++++++++++++++- mooncake-wheel/tests/test_mooncake_config.py | 116 +++++++++++++++++++ 2 files changed, 225 insertions(+), 2 deletions(-) diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/mooncake-wheel/mooncake/mooncake_config.py index 8a7b303756..8762a3445e 100644 --- a/mooncake-wheel/mooncake/mooncake_config.py +++ b/mooncake-wheel/mooncake/mooncake_config.py @@ -20,6 +20,7 @@ Transfer Engine C++ Level (Advanced): ------------------------------------- In addition to tcp and rdma, the C++ Transfer Engine also supports: +- efa: AWS Elastic Fabric Adapter transport (libfabric-based) - nvmeof: NVMe over Fabric for direct NVMe storage access - nvlink: NVIDIA NVLink for inter-GPU communication across nodes - nvlink_intra: NVIDIA NVLink for intra-node GPU communication @@ -28,6 +29,19 @@ - cxl: Compute Express Link for memory pooling and sharing - ascend: Huawei Ascend NPU communication (HCCL and direct transport) +Store Surface (mooncake-store): +------------------------------- +MooncakeConfig also drives the Mooncake Store, which additionally accepts: +- ub / ubshmem: Unified Bus transport and its shared-memory variant +- maca: MetaX MACA GPU transport +- sunrise_link: SunriseLink interconnect transport +- rpc_only: Store-only mode with no Transfer Engine attached + +Protocol names are matched case-sensitively by the C++ engine, so the value is +normalised to lowercase when a MooncakeConfig is constructed (e.g. "RDMA" -> +"rdma"). Because the authoritative, build-flag-dependent set lives in C++, an +unrecognised protocol is passed through with a warning rather than rejected. + For most use cases, 'tcp' or 'rdma' is recommended. The default is 'tcp'. For RDMA, you also need to specify the device_name (e.g., 'mlx5_0', 'erdma_0') or use auto-discovery. @@ -46,10 +60,13 @@ export MOONCAKE_DEVICE="auto-discovery" """ import json +import logging import os from dataclasses import dataclass from typing import Optional +logger = logging.getLogger(__name__) + DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 # 3.125 GiB DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 # 1.0 GiB @@ -65,6 +82,42 @@ ("b", 1), ] +# Protocols Mooncake is known to accept. This is a *diagnostic hint*, not an +# authoritative gate: MooncakeConfig drives both the Transfer Engine and the +# Store, and the C++ layer is the source of truth for what a given build +# actually supports (several transports are gated behind USE_* build flags). +# The C++ comparisons are exact and lowercase (multi_transport.cpp +# installTransport, mooncake-store client_service.cpp), so the protocol is +# canonicalised to lowercase below and an unrecognised value only warns. +# Union of Transfer Engine transports and Store-only modes. Keep in sync with +# the protocol list documented in the module docstring above. +_KNOWN_PROTOCOLS = frozenset({ + # Transfer Engine transports (mooncake-transfer-engine installTransport) + "tcp", + "rdma", + "efa", + "nvmeof", + "nvlink", + "nvlink_intra", + "hip", + "barex", + "cxl", + "ascend", + "ub", + "ubshmem", + "maca", + "sunrise_link", + # Store-only mode (mooncake-store client_service.cpp: no transfer engine) + "rpc_only", +}) + +# Required fields that must be present AND non-empty. +_REQUIRED_NON_EMPTY_FIELDS = ( + "local_hostname", + "metadata_server", + "master_server_address", +) + def _parse_segment_size(value) -> int: if isinstance(value, int): @@ -119,10 +172,14 @@ class MooncakeConfig: metadata_server (str): The address of the metadata server. global_segment_size (int): The size of each global segment in bytes. local_buffer_size (int): The size of the local buffer in bytes. - protocol (str): The communication protocol to use. Supported values: + protocol (str): The communication protocol to use. Common values: - "tcp" (default): Standard TCP/IP protocol - "rdma": RDMA protocol (requires RDMA-capable NICs and device_name) - See module docstring for full list of supported protocols. + The value is normalised to lowercase (the C++ engine matches + protocol names case-sensitively). The Transfer Engine and Store + together accept a wider set; see the module docstring. An + unrecognised value is passed through to the engine with a warning, + not rejected. device_name (Optional[str]): The name of the RDMA device to use (e.g., "mlx5_0", "erdma_0", or "auto-discovery"). Required when protocol is "rdma", optional for other protocols. @@ -170,6 +227,56 @@ class MooncakeConfig: ssd_offload_path: str = "" tenant_id: str = "default" + def __post_init__(self): + """Validate and normalise configuration invariants. + + This runs for every ``MooncakeConfig`` instance regardless of how it is + constructed (``from_file``, ``load_from_env`` or direct instantiation). + The protocol is canonicalised to lowercase (the C++ engine is + case-sensitive) and an unrecognised protocol is warned about but passed + through, because the authoritative set is decided in the + build-flag-dependent C++ layer. Genuine structural problems (a + non-string or empty protocol, a negative size, an empty required field) + are still reported here with an actionable message instead of surfacing + as a cryptic failure deep inside the C++ engine. + """ + if not isinstance(self.protocol, str) or not self.protocol.strip(): + raise ValueError( + f"Invalid protocol: {self.protocol!r}. Protocol must be a " + f"non-empty string, e.g. 'tcp' or 'rdma'." + ) + # Canonicalise to lowercase. The C++ engine matches protocol names + # case-sensitively against lowercase literals, so e.g. "RDMA" would be + # silently accepted here but rejected deep in the engine; normalising + # turns that hidden misconfiguration into a working configuration. + self.protocol = self.protocol.strip().lower() + # Warn (do not reject) on values we do not recognise: MooncakeConfig + # drives both Transfer Engine and Store paths, and a given build may + # support protocols beyond this list. The C++ layer is authoritative and + # will reject a genuinely unsupported protocol with its own error. + if self.protocol not in _KNOWN_PROTOCOLS: + logger.warning( + "Unrecognised protocol %r; passing it through to the Mooncake " + "engine unchanged. Known protocols are: %s.", + self.protocol, + ", ".join(sorted(_KNOWN_PROTOCOLS)), + ) + + for field_name in ("global_segment_size", "local_buffer_size"): + value = getattr(self, field_name) + if value < 0: + raise ValueError( + f"Invalid {field_name}: {value}. Size must be non-negative." + ) + + for field_name in _REQUIRED_NON_EMPTY_FIELDS: + value = getattr(self, field_name) + if not isinstance(value, str) or not value.strip(): + raise ValueError( + f"Config field {field_name!r} must be a non-empty string, " + f"got {value!r}." + ) + @staticmethod def from_file(file_path: str) -> 'MooncakeConfig': """Load the config from a JSON file.""" diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/mooncake-wheel/tests/test_mooncake_config.py index 5172682205..a1c91abeb6 100644 --- a/mooncake-wheel/tests/test_mooncake_config.py +++ b/mooncake-wheel/tests/test_mooncake_config.py @@ -3,6 +3,7 @@ import tempfile import unittest +from mooncake import mooncake_config as _cfg_mod from mooncake.mooncake_config import ( MooncakeConfig, DEFAULT_GLOBAL_SEGMENT_SIZE, @@ -318,5 +319,120 @@ def test_whitespace_handling(self): self.assertEqual(_parse_segment_size(" 3 gb "), 3 * 1024 ** 3) +class TestMooncakeConfigValidation(unittest.TestCase): + """Tests for the field validation performed in MooncakeConfig.__post_init__.""" + + BASE_KWARGS = dict( + local_hostname="localhost", + metadata_server="localhost:8080", + global_segment_size=DEFAULT_GLOBAL_SEGMENT_SIZE, + local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE, + protocol="tcp", + device_name="", + master_server_address="localhost:8081", + ) + + def make(self, **overrides): + kwargs = dict(self.BASE_KWARGS) + kwargs.update(overrides) + return MooncakeConfig(**kwargs) + + def test_known_protocols_normalized_to_lowercase(self): + # Mixed-case input is canonicalised to lowercase so it matches the + # case-sensitive C++ engine; surrounding whitespace is trimmed. Known + # protocols (including Store modes the old hard allowlist rejected) do + # not warn. + cases = { + "tcp": "tcp", + "rdma": "rdma", + "efa": "efa", + "RDMA": "rdma", + "Tcp": "tcp", + " rdma ": "rdma", + "cxl": "cxl", + "ascend": "ascend", + "nvlink_intra": "nvlink_intra", + "ub": "ub", + "ubshmem": "ubshmem", + "maca": "maca", + "sunrise_link": "sunrise_link", + "rpc_only": "rpc_only", + } + for given, expected in cases.items(): + with self.subTest(protocol=given): + config = self.make(protocol=given) + self.assertEqual(config.protocol, expected) + + def test_empty_protocol_raises(self): + # An empty or whitespace-only protocol is a caller error and still + # fails fast. + for protocol in ["", " "]: + with self.subTest(protocol=protocol): + with self.assertRaises(ValueError) as cm: + self.make(protocol=protocol) + self.assertIn("Invalid protocol", str(cm.exception)) + + def test_unknown_protocol_warns_but_is_passed_through(self): + # Unknown-but-non-empty values are no longer rejected in Python: + # MooncakeConfig drives both Transfer Engine and Store paths and the C++ + # layer is the source of truth. We warn and pass the lowercased value + # through. + for given, expected in [("rmda", "rmda"), ("udp", "udp"), ("Foo", "foo")]: + with self.subTest(protocol=given): + with self.assertLogs(_cfg_mod.logger, level="WARNING") as cm: + config = self.make(protocol=given) + self.assertEqual(config.protocol, expected) + self.assertTrue( + any("Unrecognised protocol" in m for m in cm.output) + ) + + def test_non_string_protocol_raises(self): + with self.assertRaises(ValueError): + self.make(protocol=None) + + def test_zero_sizes_allowed(self): + # 0 is a documented sentinel (e.g. global_segment_size == 0 disables the + # store), so it must remain valid. + config = self.make(global_segment_size=0, local_buffer_size=0) + self.assertEqual(config.global_segment_size, 0) + self.assertEqual(config.local_buffer_size, 0) + + def test_negative_sizes_raise(self): + with self.assertRaises(ValueError) as cm: + self.make(global_segment_size=-1) + self.assertIn("global_segment_size", str(cm.exception)) + with self.assertRaises(ValueError) as cm: + self.make(local_buffer_size=-1024) + self.assertIn("local_buffer_size", str(cm.exception)) + + def test_empty_required_field_raises(self): + for field in ["local_hostname", "metadata_server", "master_server_address"]: + for bad in ["", " "]: + with self.subTest(field=field, value=bad): + with self.assertRaises(ValueError) as cm: + self.make(**{field: bad}) + self.assertIn(field, str(cm.exception)) + + def test_from_file_warns_on_unknown_protocol(self): + with open(self.config_path, "w") as f: + json.dump({ + "local_hostname": "localhost", + "metadata_server": "localhost:8080", + "master_server_address": "localhost:8081", + "protocol": "rmda", # typo -> unknown, warned not rejected + }, f) + with self.assertLogs(_cfg_mod.logger, level="WARNING") as cm: + config = MooncakeConfig.from_file(self.config_path) + self.assertEqual(config.protocol, "rmda") + self.assertTrue(any("Unrecognised protocol" in m for m in cm.output)) + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.config_path = os.path.join(self._tmp.name, "config.json") + + def tearDown(self): + self._tmp.cleanup() + + if __name__ == '__main__': unittest.main() From 8ceac2732688850b17f3ed21c52ab2957f0e541e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=BA=E7=94=9F=E8=8B=A5=E5=8F=AA=E5=A6=82=E5=88=9D?= =?UTF-8?q?=E8=A7=81?= Date: Tue, 7 Jul 2026 10:35:11 +0800 Subject: [PATCH 24/57] [Store] Externalize S3 client config via environment variables (#2649) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 郭祥 Co-authored-by: Claude --- mooncake-common/include/environ.h | 35 +++++++ mooncake-common/src/environ.cpp | 37 +++++++ mooncake-common/tests/environ_test.cpp | 77 ++++++++++++++ mooncake-store/src/utils/s3_helper.cpp | 133 +++++++++---------------- 4 files changed, 197 insertions(+), 85 deletions(-) diff --git a/mooncake-common/include/environ.h b/mooncake-common/include/environ.h index 78f8f45028..ec860500ae 100644 --- a/mooncake-common/include/environ.h +++ b/mooncake-common/include/environ.h @@ -52,8 +52,30 @@ class Environ { bool GetWithNvidiaPeermem() const { return with_nvidia_peermem_; } int GetEfaCqThreads() const { return efa_cq_threads_; } + // AWS / S3 client configuration + std::string GetAwsRegion() const { return aws_region_; } + std::string GetAwsS3Endpoint() const { return aws_s3_endpoint_; } + std::string GetAwsBucketName() const { return aws_bucket_name_; } + std::string GetAwsAccessKeyId() const { return aws_access_key_id_; } + std::string GetAwsSecretAccessKey() const { return aws_secret_access_key_; } + bool GetAwsUseVirtualAddressing() const { + return aws_use_virtual_addressing_; + } + bool GetAwsUseHttps() const { return aws_use_https_; } + // Empty string means "unset" — s3_helper keeps the AWS SDK default in + // that case. Parsing to AWS enums is done by the consumer. + std::string GetAwsRequestChecksumCalculation() const { + return aws_request_checksum_calculation_; + } + std::string GetAwsResponseChecksumValidation() const { + return aws_response_checksum_validation_; + } + int64_t GetAwsConnectTimeoutMs() const { return aws_connect_timeout_ms_; } + int64_t GetAwsRequestTimeoutMs() const { return aws_request_timeout_ms_; } + // Helper method to get int from env static int GetInt(const char* name, int default_value); + static int64_t GetInt64(const char* name, int64_t default_value); // Helper method to get size_t from env static size_t GetSizeT(const char* name, size_t default_value); // Helper method to get bool from env (checks for "1", "true", "TRUE") @@ -103,6 +125,19 @@ class Environ { bool path_roundrobin_; bool with_nvidia_peermem_; int efa_cq_threads_; + + // AWS / S3 client configuration + std::string aws_region_; + std::string aws_s3_endpoint_; + std::string aws_bucket_name_; + std::string aws_access_key_id_; + std::string aws_secret_access_key_; + bool aws_use_virtual_addressing_; + bool aws_use_https_; + std::string aws_request_checksum_calculation_; + std::string aws_response_checksum_validation_; + int64_t aws_connect_timeout_ms_; + int64_t aws_request_timeout_ms_; }; } // namespace mooncake diff --git a/mooncake-common/src/environ.cpp b/mooncake-common/src/environ.cpp index 46aa1e9a9c..a920414fd9 100644 --- a/mooncake-common/src/environ.cpp +++ b/mooncake-common/src/environ.cpp @@ -31,6 +31,23 @@ int Environ::GetInt(const char* name, int default_value) { return default_value; } +int64_t Environ::GetInt64(const char* name, int64_t default_value) { + const char* val = std::getenv(name); + if (val) { + char* endptr = nullptr; + errno = 0; + long long result = std::strtoll(val, &endptr, 10); + if (endptr == val || *endptr != '\0' || errno == ERANGE) { + std::cerr << "[Mooncake] Warning: invalid value '" << val + << "' for env " << name << ", using default " + << default_value << std::endl; + return default_value; + } + return static_cast(result); + } + return default_value; +} + size_t Environ::GetSizeT(const char* name, size_t default_value) { const char* val = std::getenv(name); if (val) { @@ -108,6 +125,26 @@ Environ::Environ() { path_roundrobin_ = GetBool("MC_PATH_ROUNDROBIN", false); with_nvidia_peermem_ = GetBool("WITH_NVIDIA_PEERMEM", true); efa_cq_threads_ = GetInt("MC_EFA_CQ_THREADS", 1); + + // AWS / S3 client configuration (consumed by s3_helper.cpp) + aws_region_ = GetString("MOONCAKE_AWS_REGION", ""); + aws_s3_endpoint_ = GetString("MOONCAKE_AWS_S3_ENDPOINT", ""); + aws_bucket_name_ = GetString("MOONCAKE_AWS_BUCKET_NAME", ""); + aws_access_key_id_ = GetString("MOONCAKE_AWS_ACCESS_KEY_ID", ""); + aws_secret_access_key_ = GetString("MOONCAKE_AWS_SECRET_ACCESS_KEY", ""); + aws_use_virtual_addressing_ = + GetBool("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", true); + aws_use_https_ = GetBool("MOONCAKE_AWS_USE_HTTPS", true); + // Empty string preserves "unset" semantics — s3_helper keeps the AWS SDK + // default in that case rather than forcing a value. + aws_request_checksum_calculation_ = + GetString("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION", ""); + aws_response_checksum_validation_ = + GetString("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION", ""); + aws_connect_timeout_ms_ = + GetInt64("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", 10000); + aws_request_timeout_ms_ = + GetInt64("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", 30000); } } // namespace mooncake diff --git a/mooncake-common/tests/environ_test.cpp b/mooncake-common/tests/environ_test.cpp index 9a0b534063..b859d59359 100644 --- a/mooncake-common/tests/environ_test.cpp +++ b/mooncake-common/tests/environ_test.cpp @@ -28,9 +28,22 @@ class EnvironTest : public ::testing::Test { void clearTestEnvVars() { unsetenv("MC_TEST_INT"); + unsetenv("MC_TEST_INT64"); unsetenv("MC_TEST_SIZET"); unsetenv("MC_TEST_BOOL"); unsetenv("MC_TEST_STRING"); + // Make sure AWS vars don't leak in from the test runner's env. + unsetenv("MOONCAKE_AWS_REGION"); + unsetenv("MOONCAKE_AWS_S3_ENDPOINT"); + unsetenv("MOONCAKE_AWS_BUCKET_NAME"); + unsetenv("MOONCAKE_AWS_ACCESS_KEY_ID"); + unsetenv("MOONCAKE_AWS_SECRET_ACCESS_KEY"); + unsetenv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING"); + unsetenv("MOONCAKE_AWS_USE_HTTPS"); + unsetenv("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION"); + unsetenv("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION"); + unsetenv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS"); + unsetenv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS"); } }; @@ -85,6 +98,70 @@ TEST_F(EnvironTest, GetIntMinValue) { EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), INT_MIN); } +// --- GetInt64 --- + +TEST_F(EnvironTest, GetInt64ValidValue) { + setenv("MC_TEST_INT64", "123456789012", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 0), 123456789012LL); +} + +TEST_F(EnvironTest, GetInt64Missing) { + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 9999), 9999); +} + +TEST_F(EnvironTest, GetInt64Empty) { + setenv("MC_TEST_INT64", "", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +TEST_F(EnvironTest, GetInt64NonNumeric) { + setenv("MC_TEST_INT64", "abc", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +TEST_F(EnvironTest, GetInt64Overflow) { + setenv("MC_TEST_INT64", "99999999999999999999999999", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +// --- AWS / S3 fields --- +// +// NOTE: Environ is a singleton whose constructor caches every value the +// first time Get() is called. So all AWS env vars must be set BEFORE the +// first Environ::Get() in this process. We therefore cover the populate +// path in a single test that takes the singleton's "first call" for +// itself; the default-path behavior is implicitly covered by Environ's +// constructor defaults (any earlier test would lock the cache to defaults +// and prevent us from observing populated values here). + +TEST_F(EnvironTest, AwsFieldsPopulateFromEnv) { + setenv("MOONCAKE_AWS_REGION", "us-east-1", 1); + setenv("MOONCAKE_AWS_S3_ENDPOINT", "https://s3.example.com", 1); + setenv("MOONCAKE_AWS_BUCKET_NAME", "my-bucket", 1); + setenv("MOONCAKE_AWS_ACCESS_KEY_ID", "AKIA-test", 1); + setenv("MOONCAKE_AWS_SECRET_ACCESS_KEY", "secret", 1); + setenv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", "0", 1); + setenv("MOONCAKE_AWS_USE_HTTPS", "0", 1); + setenv("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION", "when_required", 1); + setenv("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION", "when_supported", 1); + setenv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", "5000", 1); + // Bogus request timeout should fall back to the registered default. + setenv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", "bogus", 1); + + const auto& e = Environ::Get(); + EXPECT_EQ(e.GetAwsRegion(), "us-east-1"); + EXPECT_EQ(e.GetAwsS3Endpoint(), "https://s3.example.com"); + EXPECT_EQ(e.GetAwsBucketName(), "my-bucket"); + EXPECT_EQ(e.GetAwsAccessKeyId(), "AKIA-test"); + EXPECT_EQ(e.GetAwsSecretAccessKey(), "secret"); + EXPECT_FALSE(e.GetAwsUseVirtualAddressing()); + EXPECT_FALSE(e.GetAwsUseHttps()); + EXPECT_EQ(e.GetAwsRequestChecksumCalculation(), "when_required"); + EXPECT_EQ(e.GetAwsResponseChecksumValidation(), "when_supported"); + EXPECT_EQ(e.GetAwsConnectTimeoutMs(), 5000); + EXPECT_EQ(e.GetAwsRequestTimeoutMs(), 30000); +} + // --- GetSizeT --- TEST_F(EnvironTest, GetSizeTValidValue) { diff --git a/mooncake-store/src/utils/s3_helper.cpp b/mooncake-store/src/utils/s3_helper.cpp index f951d4bbee..ce394a9904 100644 --- a/mooncake-store/src/utils/s3_helper.cpp +++ b/mooncake-store/src/utils/s3_helper.cpp @@ -10,11 +10,11 @@ #include #include #include +#include #include #include #include #include -#include #include #include #include @@ -27,69 +27,30 @@ #include #include #include -#include "utils/type_util.h" +#include +#include "environ.h" #include "fmt/format.h" namespace mooncake { namespace { -constexpr int64_t kDefaultS3ConnectTimeoutMs = 10000; -constexpr int64_t kDefaultS3RequestTimeoutMs = 30000; - -struct S3Env { - std::string region; - - std::string endpoint; - - std::string bucket; - - std::string access_key; - - std::string secret_key; - - bool use_virtual_addressing = true; - - int64_t connect_timeout_ms = kDefaultS3ConnectTimeoutMs; - - int64_t request_timeout_ms = kDefaultS3RequestTimeoutMs; -}; - -S3Env s3_env; - -void AssignStringFromEnv(const char *env_name, std::string &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - target = env_value; - } else { - target.clear(); - } -} - -void AssignBoolFromEnv(const char *env_name, bool &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - bool parsed = true; - if (TypeUtil::ParseBool(env_value, parsed)) { - target = parsed; - return; - } - LOG(WARNING) << "Invalid " << env_name << " value: " << env_value; - } -} - -void AssignTimeoutFromEnv(const char *env_name, int64_t default_value, - int64_t &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - int64_t parsed; - if (TypeUtil::ParseInt64(env_value, parsed)) { - target = parsed; - return; - } - LOG(WARNING) << "Invalid " << env_name << " value: " << env_value; - } - target = default_value; +// Parse a checksum-mode string ("when_supported" | "when_required") into the +// AWS SDK enum. Returns std::nullopt when the value is unset or invalid — the +// caller should then leave the AWS SDK's default in place. This is purely +// string-to-enum logic (not getenv), so it lives next to the AWS types it +// produces. +template +std::optional ParseChecksumMode(const std::string &value) { + if (value.empty()) return std::nullopt; + std::string lower = value; + std::transform(lower.begin(), lower.end(), lower.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (lower == "when_required") return T::WHEN_REQUIRED; + if (lower == "when_supported") return T::WHEN_SUPPORTED; + LOG(WARNING) << "Invalid value: " << value + << ", ignoring (keeping AWS SDK default)"; + return std::nullopt; } } // namespace @@ -102,21 +63,10 @@ void S3Helper::InitAPI() { Aws::InitAPI(options_); aws_initialized = true; - // Read environment variables once during initialization (fallback as needed - // if not set) - AssignStringFromEnv("MOONCAKE_AWS_REGION", s3_env.region); - AssignStringFromEnv("MOONCAKE_AWS_S3_ENDPOINT", s3_env.endpoint); - AssignStringFromEnv("MOONCAKE_AWS_BUCKET_NAME", s3_env.bucket); - AssignStringFromEnv("MOONCAKE_AWS_ACCESS_KEY_ID", s3_env.access_key); - AssignStringFromEnv("MOONCAKE_AWS_SECRET_ACCESS_KEY", s3_env.secret_key); - - AssignBoolFromEnv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", - s3_env.use_virtual_addressing); - - AssignTimeoutFromEnv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", - kDefaultS3ConnectTimeoutMs, s3_env.connect_timeout_ms); - AssignTimeoutFromEnv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", - kDefaultS3RequestTimeoutMs, s3_env.request_timeout_ms); + // Force Environ initialization so all MOONCAKE_AWS_* env vars are read + // once during startup, matching the previous "read once at InitAPI" + // behavior. + (void)Environ::Get(); } void S3Helper::ShutdownAPI() { @@ -128,30 +78,43 @@ void S3Helper::ShutdownAPI() { S3Helper::S3Helper(const std::string &endpoint, const std::string &bucket, const std::string ®ion) { + const auto &env = Environ::Get(); Aws::Client::ClientConfiguration config(true); - config.connectTimeoutMs = s3_env.connect_timeout_ms; - config.requestTimeoutMs = s3_env.request_timeout_ms; - config.scheme = Aws::Http::Scheme::HTTPS; + config.connectTimeoutMs = env.GetAwsConnectTimeoutMs(); + config.requestTimeoutMs = env.GetAwsRequestTimeoutMs(); + config.scheme = env.GetAwsUseHttps() ? Aws::Http::Scheme::HTTPS + : Aws::Http::Scheme::HTTP; + if (auto request_checksum = + ParseChecksumMode( + env.GetAwsRequestChecksumCalculation())) { + config.checksumConfig.requestChecksumCalculation = *request_checksum; + } + if (auto response_checksum = + ParseChecksumMode( + env.GetAwsResponseChecksumValidation())) { + config.checksumConfig.responseChecksumValidation = *response_checksum; + } if (!region.empty()) { config.region = region; } else { - config.region = s3_env.region; + config.region = env.GetAwsRegion(); } if (!endpoint.empty()) { config.endpointOverride = endpoint; } else { - config.endpointOverride = s3_env.endpoint; + config.endpointOverride = env.GetAwsS3Endpoint(); } - bucket_ = s3_env.bucket; + bucket_ = env.GetAwsBucketName(); if (!bucket.empty()) { bucket_ = bucket; } - Aws::Auth::AWSCredentials credentials(s3_env.access_key, s3_env.secret_key); + Aws::Auth::AWSCredentials credentials(env.GetAwsAccessKeyId(), + env.GetAwsSecretAccessKey()); // Concatenate log information into member variable connection_info_ connection_info_ = fmt::format( @@ -169,14 +132,14 @@ S3Helper::S3Helper(const std::string &endpoint, const std::string &bucket, bucket_.empty() ? "unset" : bucket_, config.connectTimeoutMs, config.requestTimeoutMs, config.scheme == Aws::Http::Scheme::HTTPS ? "HTTPS" : "HTTP", - !s3_env.access_key.empty() ? "set" : "unset", - !s3_env.secret_key.empty() ? "set" : "unset", - s3_env.use_virtual_addressing ? "true" : "false"); + !env.GetAwsAccessKeyId().empty() ? "set" : "unset", + !env.GetAwsSecretAccessKey().empty() ? "set" : "unset", + env.GetAwsUseVirtualAddressing() ? "true" : "false"); s3_client_ = Aws::S3::S3Client( credentials, config, Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never, - s3_env.use_virtual_addressing); + env.GetAwsUseVirtualAddressing()); } S3Helper::~S3Helper() = default; @@ -872,4 +835,4 @@ tl::expected S3Helper::DeleteObjectsWithPrefix( return DeleteObjects(object_keys); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake From a3585aa77cd6e791f15ca36fb77671ffd0f03dfe Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Tue, 7 Jul 2026 15:19:19 +0800 Subject: [PATCH 25/57] [TENT] Per-pool QP allocation with per-pool SL/TC (RFC #2568 step 2) (#2759) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Per-pool QP allocation with per-pool SL/TC (RFC #2568 step 2) Step 1 (#2640) added the SelectionPolicy `qp_pool` schema (parsed and carried, but not yet wired). This step makes named QP pools real at the RDMA link layer: each pool gets its own contiguous run of data QPs inside every endpoint, handshaked with that pool's Service Level / Traffic Class. This is the "allocate multiple QPs, assign different params during their handshakes" mechanism @alogfans described on the RFC. - EndPointParams::qp_pools: optional per-pool layout (name, num_qp, SL, TC). Empty (default) = today's single homogeneous run of qp_mul_factor QPs. - construct(): total QP count = sum of per-pool num_qp when pools are set; qp_pool_segments_ records each pool's [begin, begin+num_qp) span. - setupOneQP(): a QP in a pool that overrides SL/TC handshakes with the pool's values; otherwise falls back to the global endpoint SL/TC. - rdma_transport: parse `transports/rdma/endpoint/qp_pools` from config. Pool SL/TC live in the RDMA config (link-layer concept); SelectionPolicy only references a pool by name. Wire-compatible: BootstrapDesc is unchanged. Both peers derive the same pool layout from the same config, so the flat qp_num list still pairs positionally. Default path (no pools configured) is byte-for-byte the prior behavior. Slice-to-pool routing (feeding SelectionResult.qp_pool into QP selection) is a follow-up so this stays reviewable in small steps. Co-Authored-By: Claude Opus 4.8 * [TENT] Route slices to their QP pool (RFC #2568 step 3) Step 2 built per-pool QP segments with per-pool SL/TC. This step wires the SelectionResult.qp_pool (from step 1's schema) all the way down so transfers actually land on their pool's QPs: SelectionResult.qp_pool -> TaskInfo.qp_pool -> SubBatch.qp_pool -> RdmaTask.qp_pool -> submitSlices() picks a QP inside that pool's segment (selectQpInPool) - selectQpInPool(): pure router. A named, known pool folds the worker-lane candidate into that pool's [begin, begin+num_qp) span; empty/unknown pool or no pools configured => unchanged global spray. Result always in range. - submitSlices(): all slices in a list share one task/pool, so it reads slice->task->qp_pool once and routes accordingly. - qp_pool is carried like the existing device_mask, through the same task/ sub-batch plumbing, so the flow mirrors code reviewers already know. Default path (no qp_pool on the policy) is unchanged: pool_name is empty, selectQpInPool returns candidate % qp_count -- the prior behavior. Tests: SelectQpInPoolTest (6 cases) covers empty/unknown/named pools, no-segments fallback, negative-candidate clamping, and in-range invariant. Co-Authored-By: Claude Opus 4.8 * [TENT] Reject QP pool layout when any pool has non-positive num_qp Addresses review on #2759: a pool with num_qp <= 0 would produce an empty or negative [begin, begin+num_qp) span and break the router. computeQpPoolSegments now rejects the whole layout (valid=false, falls back to the default single homogeneous pool) instead of building a broken one. Adds zero/negative unit tests; qp_pool_layout_test passes 12/12. Co-Authored-By: Claude Opus 4.8 * ci: retrigger flaky build (3.10); no code change --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 --- .../tent/runtime/transfer_engine_impl.h | 1 + .../tent/include/tent/runtime/transport.h | 4 + .../include/tent/transport/rdma/endpoint.h | 8 + .../tent/include/tent/transport/rdma/params.h | 89 ++++++++ .../tent/include/tent/transport/rdma/slice.h | 5 + .../tent/src/runtime/transfer_engine_impl.cpp | 10 +- .../tent/src/transport/rdma/endpoint.cpp | 59 ++++- .../src/transport/rdma/rdma_transport.cpp | 37 +++ .../tent/tests/CMakeLists.txt | 7 + .../tent/tests/qp_pool_layout_test.cpp | 212 ++++++++++++++++++ 10 files changed, 425 insertions(+), 7 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 41aca93ad3..1bd440545f 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -55,6 +55,7 @@ struct TaskInfo { int xport_priority{0}; // transport priority (for fallback) int failover_count{0}; // number of failover attempts uint64_t device_mask{~0ULL}; // Device mask for quota allocation + std::string qp_pool; // Named QP pool (RFC #2568 step 3), "" = none Request request; bool staging{false}; TransferStatusEnum status{TransferStatusEnum::PENDING}; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index b28cf7cf2a..7f1a2d3c2b 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -55,6 +55,10 @@ class Transport { } uint64_t device_mask; // Device mask for transport selection + // Named QP pool for this batch's transfers (RFC #2568 step 3). Empty = + // no pool (default spray). Carried like device_mask, from the matched + // SelectionPolicy down to each RdmaTask. + std::string qp_pool; BatchID progress_batch_id{0}; std::function notify_progress; }; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h index 245715871c..d3c8efc1f1 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h @@ -172,6 +172,10 @@ class RdmaEndPoint : public std::enable_shared_from_this { int setupOneQP(int qp_index, const std::string& peer_gid, uint16_t peer_lid, uint32_t peer_qp_num, std::string* reply_msg = nullptr); + // Returns the pool segment owning qp_index, or nullptr when no pools are + // configured (the default single-pool case). Read-only after construct(). + const QpPoolSegment* poolForQp(int qp_index) const; + bool reserveQuota(int qp_index, int num_entries); void cancelQuota(int qp_index, int num_entries); @@ -193,6 +197,10 @@ class RdmaEndPoint : public std::enable_shared_from_this { std::string endpoint_name_; std::vector qp_list_; + // Per-pool QP layout, resolved once in construct() from params_->qp_pools. + // Empty = default single pool spanning all of qp_list_. Each segment's + // [begin, begin+num_qp) indexes into qp_list_. Read-only after construct(). + std::vector qp_pool_segments_; // Each data QP queue is owned by exactly one worker lane; reset/deconstruct // are synchronized by the endpoint lifecycle lock. std::vector slice_queue_; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h index f67906c235..b438281dd6 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h @@ -15,8 +15,12 @@ #ifndef TENT_PARAMS_H #define TENT_PARAMS_H +#include + #include #include +#include +#include namespace mooncake { namespace tent { @@ -29,6 +33,24 @@ struct DeviceParams { int max_cqe = 4096; }; +// One named QP pool. When SelectionPolicy entries declare a `qp_pool`, each +// distinct pool gets its own contiguous run of data QPs inside every endpoint, +// handshaked with that pool's SL/TC. Transfers routed to a pool only use its +// QPs, giving link-layer isolation between traffic classes (RFC #2568 step 2). +// +// Layering note: the wire handshake is unchanged. Both peers derive the same +// pool layout from the same SelectionPolicy config, so the flat qp_num list is +// still paired positionally — pool P's i-th QP on one side lines up with pool +// P's i-th QP on the other. This keeps BootstrapDesc byte-compatible with peers +// that don't know about pools (they simply run a single default pool). +struct QpPoolSegment { + std::string name; + int num_qp = 0; // Number of data QPs dedicated to this pool. + int begin = 0; // Index into qp_list_ where this pool's QPs start. + int service_level = -1; // -1 = fall back to EndPointParams::service_level. + int traffic_class = -1; // -1 = fall back to EndPointParams::traffic_class. +}; + struct EndPointParams { int endpoint_store_cap = 65536; int qp_mul_factor = 6; // Derived from RdmaParams::num_lanes. @@ -37,6 +59,12 @@ struct EndPointParams { int max_inline_bytes = 64; ibv_mtu path_mtu = IBV_MTU_4096; + // Named QP pools. Empty (default) = today's behavior: a single homogeneous + // run of qp_mul_factor data QPs, all handshaked with the global SL/TC. When + // non-empty, the segments partition the data QPs by pool; total QP count is + // the sum of per-pool num_qp. Both peers must derive the same layout. + std::vector qp_pools; + // Advanced parameters, do not change unless you understand them // INIT State uint16_t pkey_index = 0; @@ -58,6 +86,67 @@ struct EndPointParams { uint8_t max_rd_atomic = 16; }; +// Result of resolving the QP-pool layout: the concrete per-pool segments (each +// with its [begin, begin+num_qp) span filled in) and the total data-QP count. +struct QpPoolLayout { + std::vector segments; + int total_qp = 0; + bool valid = false; // false = invalid config (non-positive total). +}; + +// Pure resolver used by RdmaEndPoint::construct(). Kept free-standing (no RDMA +// handles) so the layout math can be unit-tested. Empty `pools` reproduces the +// historical single homogeneous run of `qp_mul_factor` data QPs (segments left +// empty, meaning "one default pool"); a non-empty config lays out one +// contiguous segment per pool and the total is the sum of per-pool num_qp. +inline QpPoolLayout computeQpPoolSegments( + const std::vector& pools, int qp_mul_factor) { + QpPoolLayout layout; + if (pools.empty()) { + layout.total_qp = qp_mul_factor; + } else { + for (const auto& pool : pools) { + // Every pool must claim at least one QP; a non-positive num_qp + // would produce an empty/negative [begin, begin+num_qp) span and + // break the router. Reject the whole layout so the caller falls + // back to the default single-pool behavior. + if (pool.num_qp <= 0) { + layout.segments.clear(); + layout.total_qp = 0; + layout.valid = false; + return layout; + } + QpPoolSegment seg = pool; + seg.begin = layout.total_qp; + layout.segments.push_back(seg); + layout.total_qp += pool.num_qp; + } + } + layout.valid = layout.total_qp > 0; + return layout; +} + +// Pure QP router used by RdmaEndPoint::submitSlices (RFC #2568 step 3). Given +// the resolved segments, the pool a transfer asked for, and a caller-provided +// candidate index (the worker lane), return the QP index the transfer should +// use. When the pool is named and found, the candidate is folded into that +// pool's [begin, begin+num_qp) span so the transfer only ever touches its +// pool's QPs. When no pool is named, or the name is unknown, or no pools are +// configured, the candidate passes through unchanged (default spray behavior). +// total_qp must be > 0; the result is always in [0, total_qp). +inline int selectQpInPool(const std::vector& segments, + const std::string& pool_name, int candidate, + int total_qp) { + if (candidate < 0) candidate = 0; + if (!pool_name.empty()) { + for (const auto& seg : segments) { + if (seg.name == pool_name && seg.num_qp > 0) + return seg.begin + (candidate % seg.num_qp); + } + } + return candidate % total_qp; +} + struct WorkerParams { int num_workers = 6; // Derived from RdmaParams::num_lanes. int max_retry_count = 8; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index f681c8d274..b27a6b7163 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -50,6 +51,10 @@ using RdmaTaskStorage = Slab; struct RdmaTask { int num_slices; Request request; + // Named QP pool this task's slices should use (RFC #2568 step 3). Empty = + // no pool selected: slices spray across all data QPs as before. Resolved + // from SelectionResult.qp_pool at task creation. + std::string qp_pool; volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; volatile int success_slices; diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index a70f3da089..3629b8e5b5 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -1485,6 +1485,7 @@ Status TransferEngineImpl::commitPreparedSubmit( prepared.submit_time; // Record start time for latency tracking task.type = owner.route.transport; task.device_mask = owner.route.device_mask; + if (owner.route.qp_pool) task.qp_pool = *owner.route.qp_pool; if (task.type == UNSPEC) { LOG(WARNING) << "Unable to find registered buffer for request: " << printRequest(merged_request); @@ -1534,6 +1535,8 @@ Status TransferEngineImpl::commitPreparedSubmit( // this batch should have the same policy) sub_batch->device_mask = batch->task_list[task_id_list[type][0]].device_mask; + sub_batch->qp_pool = + batch->task_list[task_id_list[type][0]].qp_pool; } auto status = transport->submitTransferTasks( @@ -1596,6 +1599,7 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, task.type = UNSPEC; task.sub_task_id = -1; task.device_mask = owner.route.device_mask; + if (owner.route.qp_pool) task.qp_pool = *owner.route.qp_pool; task.derived = task_plan.task_id != owner.owner_task_id; } @@ -1674,6 +1678,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { auto route = resolveTransport(task.request, 0); task.type = route.transport; task.device_mask = route.device_mask; + if (route.qp_pool) task.qp_pool = *route.qp_pool; if (task.type == UNSPEC) { return finishQueuedOwner(owner_id, FAILED); } @@ -1702,7 +1707,10 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { auto& transport = transport_list_[task.type]; if (!transport) return finishQueuedOwner(owner_id, FAILED); auto& sub_batch = batch->sub_batch[task.type]; - if (task.type == RDMA) sub_batch->device_mask = task.device_mask; + if (task.type == RDMA) { + sub_batch->device_mask = task.device_mask; + sub_batch->qp_pool = task.qp_pool; + } task.sub_task_id = sub_batch->size(); auto status = transport->submitTransferTasks(sub_batch, {task.request}); if (!status.ok()) { diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp index d1f5109f7d..81fac149d1 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp @@ -83,12 +83,28 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, params_ = params; endpoint_name_ = endpoint_name; inflight_slices_ = 0; - qp_list_.resize(params_->qp_mul_factor); + + // Resolve the per-pool QP layout (see computeQpPoolSegments). Empty + // qp_pools (the default) keeps the historical single homogeneous run of + // qp_mul_factor data QPs; a non-empty config lays out one contiguous + // segment per pool. qp_pool_segments_ is read-only after construct(). + QpPoolLayout layout = + computeQpPoolSegments(params_->qp_pools, params_->qp_mul_factor); + if (!layout.valid) { + LOG(ERROR) << "Invalid QP count " << layout.total_qp + << " (qp_mul_factor=" << params_->qp_mul_factor + << ", pools=" << params_->qp_pools.size() << ")"; + return -1; + } + qp_pool_segments_ = std::move(layout.segments); + const int total_qp = layout.total_qp; + + qp_list_.resize(total_qp); // Value-initialize the full array because cleanup may run after only a // prefix of QPs has been created. - wr_depth_list_ = new WrDepthBlock[params_->qp_mul_factor](); + wr_depth_list_ = new WrDepthBlock[total_qp](); - for (int i = 0; i < params_->qp_mul_factor; ++i) { + for (int i = 0; i < total_qp; ++i) { wr_depth_list_[i].value = 0; ibv_qp_init_attr attr; memset(&attr, 0, sizeof(attr)); @@ -648,6 +664,14 @@ int RdmaEndPoint::resetConnection(const std::string& reason) { return 0; } +const QpPoolSegment* RdmaEndPoint::poolForQp(int qp_index) const { + for (const auto& seg : qp_pool_segments_) { + if (qp_index >= seg.begin && qp_index < seg.begin + seg.num_qp) + return &seg; + } + return nullptr; +} + int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid, std::vector peer_qp_num_list, std::string* reply_msg) { @@ -694,7 +718,17 @@ int RdmaEndPoint::submitSlices(std::vector& slice_list, RWSpinlock::ReadGuard guard(lock_); if (qp_list_.empty()) return 0; if (qp_index < 0) qp_index = 0; - qp_index %= qp_list_.size(); + // Route to the QP pool this transfer asked for (RFC #2568 step 3). All + // slices in a list belong to one task, hence one pool; fold the worker-lane + // candidate into that pool's QP segment. Empty/unknown pool or no pools + // configured => unchanged global spray. + static const std::string kNoPool; + const std::string& pool_name = + (!slice_list.empty() && slice_list.front()->task) + ? slice_list.front()->task->qp_pool + : kNoPool; + qp_index = selectQpInPool(qp_pool_segments_, pool_name, qp_index, + (int)qp_list_.size()); // Check endpoint status before submitting if (status_.load(std::memory_order_relaxed) != EP_READY) return 0; auto cq = context_->cq(qp_index % context_->cqCount()); @@ -847,6 +881,19 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid, assert(qp_index >= 0 && qp_index < (int)qp_list_.size()); auto& qp = qp_list_[qp_index]; + // Resolve link-layer QoS for this QP. When it belongs to a pool that + // overrides SL/TC, use the pool's values; otherwise fall back to the global + // endpoint SL/TC (unchanged default behavior). + const QpPoolSegment* pool = poolForQp(qp_index); + const uint8_t qp_service_level = + (pool && pool->service_level >= 0) + ? static_cast(pool->service_level) + : params_->service_level; + const uint8_t qp_traffic_class = + (pool && pool->traffic_class >= 0) + ? static_cast(pool->traffic_class) + : params_->traffic_class; + // RESET -> INIT ibv_qp_attr attr; memset(&attr, 0, sizeof(attr)); @@ -886,9 +933,9 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid, attr.ah_attr.grh.sgid_index = context().gidIndex(); attr.ah_attr.grh.hop_limit = params_->hop_limit; attr.ah_attr.grh.flow_label = params_->flow_label; - attr.ah_attr.grh.traffic_class = params_->traffic_class; + attr.ah_attr.grh.traffic_class = qp_traffic_class; attr.ah_attr.dlid = peer_lid; - attr.ah_attr.sl = params_->service_level; + attr.ah_attr.sl = qp_service_level; attr.ah_attr.src_path_bits = params_->src_path_bits; attr.ah_attr.static_rate = params_->static_rate; attr.ah_attr.is_global = 1; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index 93736e24dd..dc7ca187a9 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -37,6 +37,7 @@ #include "tent/common/utils/string_builder.h" #include "tent/runtime/topology.h" #include "tent/common/utils/random.h" +#include "tent/thirdparty/nlohmann/json.h" #define SET_DEVICE(key, param) \ param = conf->get("transports/rdma/device/" #key, param) @@ -192,6 +193,41 @@ static Status convertConfToRdmaParams(std::shared_ptr conf, else params->endpoint.path_mtu = IBV_MTU_512; + // Optional per-pool QP layout (RFC #2568 step 2). Each entry defines a + // named pool with its own QP count and link-layer SL/TC; + // SelectionPolicy.qp_pool references these by name. Absent/empty => single + // default pool (unchanged). The pool SL/TC live here in the RDMA config, + // not in SelectionPolicy, to keep the link-layer QoS definition in the + // transport layer; policies only reference a pool by name. + params->endpoint.qp_pools.clear(); + auto qp_pools_json = + conf->getArray("transports/rdma/endpoint/qp_pools"); + for (const auto& pool_json : qp_pools_json) { + if (!pool_json.is_object()) { + LOG(WARNING) << "Ignore non-object entry in qp_pools"; + continue; + } + if (!pool_json.contains("name") || !pool_json["name"].is_string()) { + LOG(WARNING) << "Ignore qp_pool entry without a string 'name'"; + continue; + } + QpPoolSegment seg; + seg.name = pool_json["name"].get(); + seg.num_qp = pool_json.value("num_qp", 0); + if (seg.num_qp <= 0) { + LOG(WARNING) << "Ignore qp_pool '" << seg.name + << "' with non-positive num_qp " << seg.num_qp; + continue; + } + seg.service_level = pool_json.value("service_level", -1); + seg.traffic_class = pool_json.value("traffic_class", -1); + params->endpoint.qp_pools.push_back(std::move(seg)); + } + if (!params->endpoint.qp_pools.empty()) { + LOG(INFO) << "Configured " << params->endpoint.qp_pools.size() + << " QP pool(s) for per-class link-layer isolation"; + } + SET_WORKERS(max_retry_count, params->workers.max_retry_count); SET_WORKERS(block_size, params->workers.block_size); SET_WORKERS(grace_period_ns, params->workers.grace_period_ns); @@ -401,6 +437,7 @@ Status RdmaTransport::submitTransferTasks( auto* task = RdmaTaskStorage::Get().allocate(); rdma_batch->task_list.push_back(task); task->request = request; + task->qp_pool = rdma_batch->qp_pool; // RFC #2568 step 3 task->num_slices = 0; task->status_word = PENDING; task->transferred_bytes = 0; diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 7329976706..868951f531 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -33,6 +33,13 @@ target_include_directories(tent_ip_utils_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_ip_utils_test COMMAND tent_ip_utils_test) +add_executable(tent_qp_pool_layout_test qp_pool_layout_test.cpp) +target_link_libraries(tent_qp_pool_layout_test PRIVATE tent_common gtest + gtest_main) +target_include_directories(tent_qp_pool_layout_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_qp_pool_layout_test COMMAND tent_qp_pool_layout_test) + add_executable(tent_coalesce_regions_test coalesce_regions_test.cpp) target_link_libraries(tent_coalesce_regions_test PRIVATE tent_common gtest gtest_main) diff --git a/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp b/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp new file mode 100644 index 0000000000..64c841c980 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp @@ -0,0 +1,212 @@ +// Copyright 2025 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Unit tests for computeQpPoolSegments — the pure QP-pool layout resolver used +// by RdmaEndPoint::construct() (RFC #2568 step 2). Kept free of RDMA handles so +// the layout math is testable without a device. + +#include + +#include "tent/transport/rdma/params.h" + +namespace mooncake { +namespace tent { +namespace { + +// Default path: no pools configured => a single homogeneous run of +// qp_mul_factor QPs, no explicit segments (poolForQp will return nullptr and +// callers fall back to the global SL/TC — byte-for-byte the prior behavior). +TEST(QpPoolLayoutTest, EmptyPoolsKeepsFlatQpMulFactor) { + auto layout = computeQpPoolSegments({}, 6); + EXPECT_TRUE(layout.valid); + EXPECT_EQ(layout.total_qp, 6); + EXPECT_TRUE(layout.segments.empty()); +} + +// A non-positive total (e.g. qp_mul_factor <= 0 with no pools) is rejected so +// construct() can fail cleanly instead of allocating a zero-length QP array. +TEST(QpPoolLayoutTest, EmptyPoolsWithNonPositiveFactorIsInvalid) { + auto layout = computeQpPoolSegments({}, 0); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); +} + +// Multiple pools lay out contiguous, non-overlapping segments; total is the +// sum of per-pool num_qp; qp_mul_factor is ignored once pools are set. +TEST(QpPoolLayoutTest, MultiplePoolsLayoutContiguousSegments) { + std::vector pools; + QpPoolSegment kv; + kv.name = "kv"; + kv.num_qp = 4; + kv.service_level = 5; + kv.traffic_class = 96; + pools.push_back(kv); + QpPoolSegment ctrl; + ctrl.name = "ctrl"; + ctrl.num_qp = 2; + pools.push_back(ctrl); + + auto layout = computeQpPoolSegments(pools, /*qp_mul_factor=*/6); + ASSERT_TRUE(layout.valid); + EXPECT_EQ(layout.total_qp, 6); // 4 + 2, not qp_mul_factor + ASSERT_EQ(layout.segments.size(), 2u); + + EXPECT_EQ(layout.segments[0].name, "kv"); + EXPECT_EQ(layout.segments[0].begin, 0); + EXPECT_EQ(layout.segments[0].num_qp, 4); + EXPECT_EQ(layout.segments[0].service_level, 5); + EXPECT_EQ(layout.segments[0].traffic_class, 96); + + EXPECT_EQ(layout.segments[1].name, "ctrl"); + EXPECT_EQ(layout.segments[1].begin, 4); // starts after kv's 4 QPs + EXPECT_EQ(layout.segments[1].num_qp, 2); + // ctrl left SL/TC unset -> sentinel -1 (setupOneQP falls back to global). + EXPECT_EQ(layout.segments[1].service_level, -1); + EXPECT_EQ(layout.segments[1].traffic_class, -1); +} + +// Segments partition [0, total_qp): every QP index maps to exactly one pool, +// mirroring RdmaEndPoint::poolForQp's linear scan. +TEST(QpPoolLayoutTest, SegmentsPartitionAllQpIndices) { + std::vector pools; + QpPoolSegment a; + a.name = "a"; + a.num_qp = 3; + pools.push_back(a); + QpPoolSegment b; + b.name = "b"; + b.num_qp = 1; + pools.push_back(b); + + auto layout = computeQpPoolSegments(pools, 6); + ASSERT_TRUE(layout.valid); + ASSERT_EQ(layout.total_qp, 4); + + auto pool_of = [&](int qp_index) -> const QpPoolSegment* { + for (const auto& seg : layout.segments) { + if (qp_index >= seg.begin && qp_index < seg.begin + seg.num_qp) + return &seg; + } + return nullptr; + }; + ASSERT_NE(pool_of(0), nullptr); + EXPECT_EQ(pool_of(0)->name, "a"); + EXPECT_EQ(pool_of(2)->name, "a"); + ASSERT_NE(pool_of(3), nullptr); + EXPECT_EQ(pool_of(3)->name, "b"); + // Out of range => no pool (default single-pool fallback in poolForQp). + EXPECT_EQ(pool_of(4), nullptr); +} + +// --- selectQpInPool: the step-3 router (slice pool -> QP index) +// --------------- + +// Helper: a two-pool layout kv=[0,4), ctrl=[4,6). +static std::vector twoPools() { + auto layout = computeQpPoolSegments( + {{"kv", 4, 0, -1, -1}, {"ctrl", 2, 0, -1, -1}}, 6); + return layout.segments; +} + +// Empty pool name => pass through, folded into the whole QP range. This is the +// default (no pool selected) behavior — identical to the pre-step-3 spray. +TEST(SelectQpInPoolTest, EmptyPoolNameSpraysAcrossAllQps) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "", 0, 6), 0); + EXPECT_EQ(selectQpInPool(segs, "", 5, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "", 7, 6), 1); // 7 % 6 +} + +// No pools configured at all => also pass through (single default pool). +TEST(SelectQpInPoolTest, NoSegmentsSpraysAcrossAllQps) { + std::vector none; + EXPECT_EQ(selectQpInPool(none, "kv", 3, 6), 3); + EXPECT_EQ(selectQpInPool(none, "", 8, 6), 2); // 8 % 6 +} + +// A named pool folds the candidate into that pool's segment only. +TEST(SelectQpInPoolTest, NamedPoolFoldsIntoItsSegment) { + auto segs = twoPools(); // kv=[0,4), ctrl=[4,6) + // kv: begin 0, num 4 -> indices 0..3 + EXPECT_EQ(selectQpInPool(segs, "kv", 0, 6), 0); + EXPECT_EQ(selectQpInPool(segs, "kv", 3, 6), 3); + EXPECT_EQ(selectQpInPool(segs, "kv", 4, 6), 0); // 4 % 4 -> begin+0 + EXPECT_EQ(selectQpInPool(segs, "kv", 6, 6), 2); // 6 % 4 -> begin+2 + // ctrl: begin 4, num 2 -> indices 4..5 + EXPECT_EQ(selectQpInPool(segs, "ctrl", 0, 6), 4); + EXPECT_EQ(selectQpInPool(segs, "ctrl", 1, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "ctrl", 3, 6), 5); // 3 % 2 -> begin+1 +} + +// Unknown pool name => fall back to the whole range (don't drop the transfer). +TEST(SelectQpInPoolTest, UnknownPoolFallsBackToWholeRange) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "nope", 5, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "nope", 9, 6), 3); // 9 % 6 +} + +// Negative candidate is clamped to 0 before folding. +TEST(SelectQpInPoolTest, NegativeCandidateClampsToZero) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "ctrl", -1, 6), 4); // begin+0 + EXPECT_EQ(selectQpInPool(segs, "", -1, 6), 0); +} + +// Every result stays in [0, total_qp) regardless of pool/candidate. +TEST(SelectQpInPoolTest, ResultAlwaysInRange) { + auto segs = twoPools(); + for (int c = 0; c < 20; ++c) { + for (const char* name : {"", "kv", "ctrl", "nope"}) { + int idx = selectQpInPool(segs, name, c, 6); + EXPECT_GE(idx, 0); + EXPECT_LT(idx, 6); + } + } +} + +// A pool with a non-positive num_qp would create an empty/negative QP span and +// break the router, so the whole layout is rejected (falls back to default). +TEST(QpPoolLayoutTest, PoolWithZeroQpIsInvalid) { + std::vector pools; + QpPoolSegment ok; + ok.name = "kv"; + ok.num_qp = 4; + pools.push_back(ok); + QpPoolSegment bad; + bad.name = "ctrl"; + bad.num_qp = 0; // invalid + pools.push_back(bad); + + auto layout = computeQpPoolSegments(pools, 6); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); + EXPECT_TRUE(layout.segments.empty()); +} + +TEST(QpPoolLayoutTest, PoolWithNegativeQpIsInvalid) { + std::vector pools; + QpPoolSegment bad; + bad.name = "kv"; + bad.num_qp = -1; // invalid + pools.push_back(bad); + + auto layout = computeQpPoolSegments(pools, 6); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); + EXPECT_TRUE(layout.segments.empty()); +} + +} // namespace +} // namespace tent +} // namespace mooncake From fbf32ca60f6bb31c96b055ad85604feb95bcbc00 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Tue, 7 Jul 2026 15:25:01 +0800 Subject: [PATCH 26/57] [TENT] Make RDMA NIC allow/deny list configurable via MC_FILTER_NIC(_EXCLUDE) (#2760) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Wire MC_TE_FILTERS(_EXCLUDE) env to the NIC allow/deny filter The platform probes already support NIC allow/deny lists via the config keys `topology/rdma_whitelist` / `topology/rdma_blacklist` (see filterInfiniBandDevices in the platform probes), but there was no environment variable to set them, so users had no way to restrict which RDMA NICs the engine discovers and uses. The legacy Transfer Engine already exposes a device whitelist via `MC_TE_FILTERS` (comma-separated). This reuses the *same* env var and semantics for the TENT path so a single variable scopes NICs across both engines, and adds `MC_TE_FILTERS_EXCLUDE` as a deny-list (the legacy engine has an allow-list only). On multi-NIC / multi-NUMA hosts this matters: the rail selection enumerates every discovered remote device and will attempt QPs to devices that are not reachable from the peer, which fail at modify-QP-to-RTR and drag down aggregate throughput (see #2758; related mapping issue #2467). This adds: - setArrayConfig(): parse a comma-separated env value into a string array (trims whitespace, drops empty items), reusing the approach of parseDoubleArray. - MC_TE_FILTERS -> topology/rdma_whitelist (allow-list, takes precedence) - MC_TE_FILTERS_EXCLUDE -> topology/rdma_blacklist (deny-list) - Unit tests for both mappings, whitespace trimming, and the unset default. - Env var docs. Opt-in: unset = discover all NICs (unchanged default behavior). Measured on an 8x H20 / CX-7 200G RoCEv2, 4-bond dual-NUMA box (2 nodes, TENT backend): without the filter, cross-node transfer hits repeated `modify QP to RTR error 22` on the NUMA1 bonds and runs at 0.14 GB/s; with `MC_TE_FILTERS=mlx5_bond_0` the topology is scoped to the reachable NIC, the errors disappear, and throughput recovers to 48 GB/s. Co-Authored-By: Claude Opus 4.8 * ci: retrigger flaky tent-ci (cuda-off); no code change --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 --- docs/source/design/transfer-engine/index.md | 2 + .../tent/src/common/config.cpp | 26 +++++++++++ .../transfer_engine_config_override_test.cpp | 43 +++++++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md index f089747be9..f80814e22d 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -471,6 +471,8 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `MC_WORKERS_PER_CTX` The number of asynchronous worker threads corresponding to each device instance - `MC_SLICE_SIZE` The segmentation granularity of user requests in Transfer Engine - `MC_RETRY_CNT` The maximum number of retries in Transfer Engine +- `MC_TE_FILTERS` Restrict which RDMA NICs the engine discovers and uses, as a comma-separated allow-list of device names (e.g. `mlx5_bond_0,mlx5_bond_1`). Only the listed NICs are kept; all others are ignored. Unset (default) discovers all NICs. This is the **same env var and semantics as the legacy Transfer Engine's device whitelist** (see below), so a single variable scopes NICs across both engines. Useful on multi-NIC / multi-NUMA hosts to keep the engine (and its rail selection) off NICs that are not routable to the peer. +- `MC_TE_FILTERS_EXCLUDE` The deny-list counterpart of `MC_TE_FILTERS`: a comma-separated list of device names to exclude from discovery. Ignored if `MC_TE_FILTERS` is set (allow-list takes precedence). Unset (default) excludes nothing. (New; the legacy engine has an allow-list only.) - `MC_AUTO_GID_MAX_RETRIES` The maximum number of automatic local GID reprobe retries during classic RDMA handshake recovery. Default value 2. Set to 0 to disable automatic GID retry. - `MC_LOG_LEVEL` This option can be set as `TRACE`/`INFO`/`WARNING`/`ERROR` (see [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)), and more detailed logs will be output during runtime - `MC_DISABLE_METACACHE` Disable local meta cache to prevent transfer failure due to dynamic memory registrations, which may downgrades the performance diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index 2a0bf8f751..a7cc9dfe8e 100644 --- a/mooncake-transfer-engine/tent/src/common/config.cpp +++ b/mooncake-transfer-engine/tent/src/common/config.cpp @@ -55,6 +55,24 @@ static inline void setConfig(Config& config, const std::string& env_key, if (val) config.setFromString(config_key, std::string(val)); } +// Like setConfig, but parses the env value as a comma-separated list and +// stores it as a string array. Empty/whitespace-only items are dropped so a +// trailing comma or spaces around names are tolerated (e.g. "mlx5_0, mlx5_1"). +static inline void setArrayConfig(Config& config, const std::string& env_key, + const std::string& config_key) { + const char* val = std::getenv(env_key.c_str()); + if (!val) return; + std::vector items; + std::stringstream ss(val); + std::string item; + while (std::getline(ss, item, ',')) { + item.erase(0, item.find_first_not_of(" \t")); + item.erase(item.find_last_not_of(" \t") + 1); + if (!item.empty()) items.push_back(item); + } + if (!items.empty()) config.set(config_key, items); +} + Status ConfigHelper::loadFromEnv(Config& config) { const char* conf_str = std::getenv("MC_TENT_CONF"); Status status = Status::OK(); @@ -119,6 +137,14 @@ Status ConfigHelper::loadFromEnv(Config& config) { "transports/rdma/disable_gpu_direct_rdma"); setConfig(config, "MC_LOG_RDMA_SLICE_AFFINITY", "transports/rdma/log_slice_affinity"); + // Restrict which RDMA NICs the engine discovers/uses (comma-separated + // device names). MC_TE_FILTERS is an allow-list — same name and semantics + // as the legacy Transfer Engine's device whitelist, so a single env works + // across both engines. MC_TE_FILTERS_EXCLUDE is a deny-list (new; the + // legacy engine has no deny-list). Unset = discover all (default). + // Consumed by filterInfiniBandDevices() in the platform probes. + setArrayConfig(config, "MC_TE_FILTERS", "topology/rdma_whitelist"); + setArrayConfig(config, "MC_TE_FILTERS_EXCLUDE", "topology/rdma_blacklist"); return status; } diff --git a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp index 223e2cf1ea..5182182762 100644 --- a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp @@ -293,6 +293,49 @@ TEST(TransferEngineConfigOverrideTest, EXPECT_TRUE(config.get("transports/rdma/log_slice_affinity", false)); } +TEST(TransferEngineConfigOverrideTest, FilterNicEnvLoadsRdmaWhitelist) { + EnvVarGuard guard("MC_TE_FILTERS", "mlx5_0,mlx5_1"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_0", "mlx5_1"}; + EXPECT_EQ(config.getArray("topology/rdma_whitelist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicExcludeEnvLoadsRdmaBlacklist) { + EnvVarGuard guard("MC_TE_FILTERS_EXCLUDE", "mlx5_2,mlx5_3"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_2", "mlx5_3"}; + EXPECT_EQ(config.getArray("topology/rdma_blacklist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicEnvTrimsWhitespaceAndEmpties) { + // Spaces around names and a trailing comma must be tolerated. + EnvVarGuard guard("MC_TE_FILTERS", " mlx5_0 , mlx5_1 ,"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_0", "mlx5_1"}; + EXPECT_EQ(config.getArray("topology/rdma_whitelist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicUnsetLeavesWhitelistEmpty) { + // Not setting the env var must leave the default (discover all NICs). + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + EXPECT_TRUE( + config.getArray("topology/rdma_whitelist").empty()); +} + TEST(TransferEngineConfigOverrideTest, ExplicitMetadataOverridesDriveSuccessfulHttpInitialization) { #ifdef _WIN32 From f9aef2ef8fb6fc517f549ea3fb3e448e127169f4 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 7 Jul 2026 06:15:38 -0400 Subject: [PATCH 27/57] [CI] Extract reusable wheel build/publish workflows (#2723) Signed-off-by: Michael Goin Co-authored-by: Claude --- .github/workflows/_build-wheel.yaml | 184 ++++++++ .github/workflows/_publish-wheel.yaml | 49 +++ .github/workflows/pre-release.yaml | 532 +++--------------------- .github/workflows/release-cuda13.yaml | 240 ++--------- .github/workflows/release-non-cuda.yaml | 126 +----- .github/workflows/release.yaml | 237 ++--------- 6 files changed, 372 insertions(+), 996 deletions(-) create mode 100644 .github/workflows/_build-wheel.yaml create mode 100644 .github/workflows/_publish-wheel.yaml diff --git a/.github/workflows/_build-wheel.yaml b/.github/workflows/_build-wheel.yaml new file mode 100644 index 0000000000..4cf5da693a --- /dev/null +++ b/.github/workflows/_build-wheel.yaml @@ -0,0 +1,184 @@ +name: _build-wheel + +# Shared build for one wheel variant (matrixed over Python version), used by the +# Release and Pre-Release workflows so both build identically. + +on: + workflow_call: + inputs: + runner: + type: string + required: true + container: + # x86 builds run in manylinux2_28 so the wheel's libstdc++/glibc floor + # stays low enough to import on RHEL/Rocky/Alma 8+. '' = bare runner. + type: string + default: '' + python-versions: + type: string + default: '["3.10", "3.11", "3.12", "3.13"]' + cuda: + # none | container | sbsa-12.8 | sbsa-13.0 + type: string + default: none + cmake-args: + # Space-separated -D flags. No semicolons; use ep-torch-versions for those. + type: string + required: true + cmake-generator: + type: string + default: '' + ep-torch-versions: + type: string + default: '' + build-with-ep: + type: string + default: '0' + variant-flag: + # build_wheel.sh variant set to 1, e.g. CU13_BUILD, NON_CUDA_BUILD + type: string + default: '' + torch-cuda-arch-list: + type: string + default: '' + build-nvlink-allocator: + type: boolean + default: false + artifact-prefix: + type: string + required: true + +env: + SCCACHE_GHA_ENABLED: "true" + +jobs: + build: + runs-on: ${{ inputs.runner }} + container: ${{ inputs.container }} + permissions: + contents: read + strategy: + matrix: + python-version: ${{ fromJSON(inputs.python-versions) }} + env: + BUILD_WITH_EP: ${{ inputs.build-with-ep }} + TORCH_CUDA_ARCH_LIST: ${{ inputs.torch-cuda-arch-list }} + CMAKE_ARGS: ${{ inputs.cmake-args }} + CMAKE_GEN: ${{ inputs.cmake-generator }} + EP_TORCH_VERSIONS_INPUT: ${{ inputs.ep-torch-versions }} + VARIANT_FLAG: ${{ inputs.variant-flag }} + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Mark workspace safe for git (container runs as root) + if: ${{ inputs.container != '' }} + run: git config --global --add safe.directory '*' + + - name: Set version from tag + run: echo "VERSION=${GITHUB_REF_NAME#v}" >> "$GITHUB_ENV" + + - name: Select Python ${{ matrix.python-version }} from manylinux image + if: ${{ inputs.container != '' }} + run: | + PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.') + PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin" + echo "$PYBIN" >> "$GITHUB_PATH" + "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel + + - name: Set up Python ${{ matrix.python-version }} + if: ${{ inputs.container == '' }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Free up disk space + if: ${{ inputs.container == '' }} + run: | + sudo rm -rf /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL /usr/local/lib/android + df -h + + - name: Install CUDA Toolkit (arm64 SBSA) + if: ${{ startsWith(inputs.cuda, 'sbsa-') }} + run: | + ver="${{ inputs.cuda }}"; ver="${ver#sbsa-}" + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb + sudo dpkg -i cuda-keyring_1.1-1_all.deb + sudo apt-get update + sudo apt-get install -y "cuda-toolkit-${ver/./-}" + echo "/usr/local/cuda/bin" >> "$GITHUB_PATH" + /usr/local/cuda/bin/nvcc --version + + - name: Run sccache-cache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Configure sccache + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Configure project + run: | + SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo" + if [ -n "$SUDO" ] && command -v apt-get >/dev/null 2>&1; then + $SUDO apt-get update -y || true + fi + $SUDO bash -x dependencies.sh -y + echo "/usr/local/go/bin" >> "$GITHUB_PATH" + gen=(); [ -n "$CMAKE_GEN" ] && gen=(-G "$CMAKE_GEN") + ep=(); [ -n "$EP_TORCH_VERSIONS_INPUT" ] && ep=(-DEP_TORCH_VERSIONS="$EP_TORCH_VERSIONS_INPUT") + mkdir -p build && cd build + # shellcheck disable=SC2086 + cmake "${gen[@]}" .. $CMAKE_ARGS "${ep[@]}" -DPython3_EXECUTABLE="$(which python3)" + + - name: Build project + run: | + for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do + [ -d "$dir" ] && export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" + done + [ -d /usr/local/cuda ] && export CUDA_HOME=/usr/local/cuda + SUDO=""; [ "$(id -u)" -ne 0 ] && SUDO="sudo -E" + cd build + cmake --build . -j"$(nproc)" + $SUDO cmake --install . + + - name: Build nvlink_allocator.so + if: ${{ inputs.build-nvlink-allocator }} + run: | + export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH + if [ -d /usr/local/cuda/lib64/stubs ]; then + export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH:-} + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LIBRARY_PATH:-} + fi + mkdir -p build/mooncake-transfer-engine/nvlink-allocator + cd mooncake-transfer-engine/nvlink-allocator + bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ + + - name: Run sccache stat for check + if: ${{ env.SCCACHE_PATH != '' }} + run: ${SCCACHE_PATH} --show-stats + + - name: Generate Python version tag + id: pytag + run: echo "tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> "$GITHUB_OUTPUT" + + - name: Build Python wheel + run: | + [ -d /usr/local/cuda ] && export CUDA_HOME=/usr/local/cuda + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" + variant=() + [ -n "$VARIANT_FLAG" ] && variant=("$VARIANT_FLAG=1") + env "${variant[@]}" \ + PYTHON_VERSION="${{ matrix.python-version }}" \ + OUTPUT_DIR="dist-py${{ steps.pytag.outputs.tag }}" \ + ./scripts/build_wheel.sh + env: + VERSION: ${{ env.VERSION }} + + - name: Upload Python wheel artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ inputs.artifact-prefix }}-py${{ steps.pytag.outputs.tag }} + path: mooncake-wheel/dist-py${{ steps.pytag.outputs.tag }}/*.whl diff --git a/.github/workflows/_publish-wheel.yaml b/.github/workflows/_publish-wheel.yaml new file mode 100644 index 0000000000..c5bab698dd --- /dev/null +++ b/.github/workflows/_publish-wheel.yaml @@ -0,0 +1,49 @@ +name: _publish-wheel + +# Shared publish tail: collect this run's wheels, attach to the GitHub Release, +# and upload to PyPI. Used by the Release workflows. + +on: + workflow_call: + inputs: + artifact-pattern: + type: string + required: true + secrets: + pypi-token: + required: false + +jobs: + publish: + runs-on: ubuntu-22.04 + permissions: + contents: write + id-token: write + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Download all wheel artifacts + uses: actions/download-artifact@v4 + with: + path: mooncake-wheel/dist-all + pattern: ${{ inputs.artifact-pattern }} + + - name: Prepare wheels for release + run: | + mkdir -p mooncake-wheel/dist-release + find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; + echo "Collected wheels for release:" + ls -la mooncake-wheel/dist-release/ + + - name: Upload wheels to GitHub Release + uses: softprops/action-gh-release@v1 + with: + files: mooncake-wheel/dist-release/*.whl + + - name: Publish package to PyPI + if: ${{ github.repository == 'kvcache-ai/Mooncake' }} + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: mooncake-wheel/dist-release/ + password: ${{ secrets.pypi-token }} diff --git a/.github/workflows/pre-release.yaml b/.github/workflows/pre-release.yaml index 7561c21b97..279d48876a 100644 --- a/.github/workflows/pre-release.yaml +++ b/.github/workflows/pre-release.yaml @@ -1,7 +1,8 @@ name: Pre-Release -# Dry-run of the release pipelines: build wheels like Release / Release Non-CUDA / -# Release CUDA 13, validate artifacts, but do not create a GitHub Release or publish to PyPI. +# Dry run of the release pipelines: build wheels through the same _build-wheel.yaml +# the Release / Release Non-CUDA / Release CUDA 13 workflows use, validate the +# artifacts, but do not create a GitHub Release or publish to PyPI. # # Trigger by pushing a pre-release tag, for example: # git tag v1.0.0-rc1 && git push origin v1.0.0-rc1 @@ -13,484 +14,75 @@ on: - 'v*-beta*' - 'v*-pre*' -env: - SCCACHE_GHA_ENABLED: "true" - jobs: build-cuda: - name: Build (CUDA 12) - runs-on: ubuntu-22.04 - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h - - - name: Install CUDA Toolkit - uses: Jimver/cuda-toolkit@v0.2.24 - with: - cuda: '12.8.1' - linux-local-args: '["--toolkit"]' - method: 'network' - sub-packages: '["nvcc", "nvrtc-dev"]' - non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - sudo -E make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '1' + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: true + cmake-args: >- + -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON + -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda build-non-cuda: - name: Build (Non-CUDA) - runs-on: ubuntu-22.04 - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "0" - NON_CUDA_BUILD: "1" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - cd build - make -j - sudo -E make install - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-non-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '0' + variant-flag: NON_CUDA_BUILD + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF + -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-non-cuda build-cuda13: - name: Build (CUDA 13) - runs-on: ubuntu-22.04 - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - sudo rm -rf /usr/local/lib/android - df -h - - - name: Install CUDA Toolkit 13 - uses: Jimver/cuda-toolkit@v0.2.29 - with: - cuda: '13.0.2' - linux-local-args: '["--toolkit"]' - method: 'network' - sub-packages: '["nvcc", "nvrtc-dev"]' - non-cuda-sub-packages: '["libcusparse-dev", "libcublas-dev", "libcusolver-dev"]' - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - sudo make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda13.0 + cuda: container + build-with-ep: '1' + variant-flag: CU13_BUILD + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: true + cmake-args: >- + -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON + -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda13 build-cuda-arm64: - name: Build (CUDA 12, arm64) - runs-on: ubuntu-22.04-arm - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Install CUDA Toolkit 12.8 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-12-8 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --install . - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_arm64 - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-12.8 + cmake-generator: Ninja + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda-arm64 build-cuda13-arm64: - name: Build (CUDA 13, arm64) - runs-on: ubuntu-22.04-arm - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Install CUDA Toolkit 13.0 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-13-0 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --install . - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_arm64 - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda13-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-13.0 + cmake-generator: Ninja + variant-flag: CU13_BUILD + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda13-arm64 validate-release: name: Validate release artifacts diff --git a/.github/workflows/release-cuda13.yaml b/.github/workflows/release-cuda13.yaml index 80d2f61192..929d9f9fa6 100644 --- a/.github/workflows/release-cuda13.yaml +++ b/.github/workflows/release-cuda13.yaml @@ -5,225 +5,45 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: build: - runs-on: ubuntu-22.04 - container: pytorch/manylinux2_28-builder:cuda13.0 - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Mark workspace safe for git (container runs as root) - run: git config --global --add safe.directory '*' - - - name: Select Python ${{ matrix.python-version }} from manylinux image - run: | - PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.') - PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin" - echo "$PYBIN" >> "$GITHUB_PATH" - "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - bash -x dependencies.sh -y - echo "/usr/local/go/bin" >> "$GITHUB_PATH" - mkdir build - cd build - cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)" - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-cuda13-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda13.0 + cuda: container + build-with-ep: '1' + variant-flag: CU13_BUILD + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: true + cmake-args: >- + -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON + -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-cuda13 build-arm64: if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04-arm - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Install CUDA Toolkit 13.0 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-13-0 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --install . - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-cuda13-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-13.0 + cmake-generator: Ninja + variant-flag: CU13_BUILD + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-cuda13-arm64 publish-release: if: ${{ !contains(github.ref_name, '-') }} needs: [build, build-arm64] - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - pattern: mooncake-wheel-cuda13* - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_CU13_API_TOKEN }} + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-cuda13*' + secrets: + pypi-token: ${{ secrets.PYPI_CU13_API_TOKEN }} diff --git a/.github/workflows/release-non-cuda.yaml b/.github/workflows/release-non-cuda.yaml index 1852717366..a3838b7937 100644 --- a/.github/workflows/release-non-cuda.yaml +++ b/.github/workflows/release-non-cuda.yaml @@ -5,120 +5,30 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: + # manylinux2_28 for the toolchain only (USE_CUDA=OFF); keeps the glibc floor + # aligned with the CUDA wheels. build: - runs-on: ubuntu-22.04 - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "0" - NON_CUDA_BUILD: "1" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - cd build - make -j - sudo make install - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-non-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '0' + variant-flag: NON_CUDA_BUILD + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF + -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-non-cuda publish-release: if: ${{ !contains(github.ref_name, '-') }} needs: build - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - pattern: mooncake-wheel-non-cuda-py* - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_API_TOKEN }} + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel-non-cuda*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index d0da9e5552..007e4a5c72 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,223 +5,44 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: + # Skip semver pre-release tags (e.g. v1.0.0-rc1); those are handled by pre-release.yaml. build: - # Skip semver pre-release tags (e.g. v1.0.0-rc1); those are handled by pre-release.yaml. if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04 - container: pytorch/manylinux2_28-builder:cuda12.8 - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - BUILD_WITH_EP: "1" - TORCH_CUDA_ARCH_LIST: "8.0;9.0" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Mark workspace safe for git (container runs as root) - run: git config --global --add safe.directory '*' - - - name: Select Python ${{ matrix.python-version }} from manylinux image - run: | - PYV_NODOT=$(echo "${{ matrix.python-version }}" | tr -d '.') - PYBIN="/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin" - echo "$PYBIN" >> "$GITHUB_PATH" - "$PYBIN/pip" install --quiet "cmake<4" setuptools wheel - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - bash -x dependencies.sh -y - echo "/usr/local/go/bin" >> "$GITHUB_PATH" - mkdir build - cd build - cmake .. -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)" - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - make install - shell: bash - - - name: Build nvlink_allocator.so - run: | - export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH - export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - mkdir -p build/mooncake-transfer-engine/nvlink-allocator - cd mooncake-transfer-engine/nvlink-allocator - bash build.sh ../../build/mooncake-transfer-engine/nvlink-allocator/ - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - # Set LD_LIBRARY_PATH for wheel building - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '1' + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: true + cmake-args: >- + -DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=ON + -DWITH_EP=ON -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel build-arm64: if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04-arm - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Free up disk space - run: | - sudo rm -rf /usr/share/dotnet - sudo rm -rf /opt/ghc - sudo rm -rf /opt/hostedtoolcache/CodeQL - - - name: Install CUDA Toolkit 12.8 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-12-8 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --install . - shell: bash - - - name: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - name: Generate Python version tag - id: generate_tag_release - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:/usr/local/lib" - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_release.outputs.python_version_tag }} ./scripts/build_wheel.sh - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-12.8 + cmake-generator: Ninja + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-arm64 publish-release: needs: [build, build-arm64] - runs-on: ubuntu-22.04 permissions: contents: write id-token: write - steps: - - name: Checkout source - uses: actions/checkout@v4 - - - name: Download all wheel artifacts - uses: actions/download-artifact@v4 - with: - path: mooncake-wheel/dist-all - - - name: Prepare wheels for release - run: | - # Move all wheels to a single directory - mkdir -p mooncake-wheel/dist-release - find mooncake-wheel/dist-all -name "*.whl" -exec cp {} mooncake-wheel/dist-release/ \; - ls -la mooncake-wheel/dist-release/ - # List all collected wheels - echo "Collected wheels for release:" - ls -la mooncake-wheel/dist-release/ - - - name: Upload wheels to GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: mooncake-wheel/dist-release/*.whl - - - name: Publish package to PyPI - if: github.repository == 'kvcache-ai/Mooncake' - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: mooncake-wheel/dist-release/ - password: ${{ secrets.PYPI_API_TOKEN }} + uses: ./.github/workflows/_publish-wheel.yaml + with: + artifact-pattern: 'mooncake-wheel*' + secrets: + pypi-token: ${{ secrets.PYPI_API_TOKEN }} From 2e21a45e8f27a3c3100609a632f131c231b3db42 Mon Sep 17 00:00:00 2001 From: tancz <544463199@qq.com> Date: Tue, 7 Jul 2026 18:16:29 +0800 Subject: [PATCH 28/57] [Store] Support local_buffer_size in mooncake_client (#2740) (#2739) Previously hardcoded to 0, now configurable with default "0" to preserve existing behavior. Signed-off-by: tan changzhi <544463199@qq.com> --- mooncake-store/src/real_client_main.cpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index fc1528b063..e912e0604d 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -17,6 +17,7 @@ DEFINE_string(master_server_address, "127.0.0.1:50051", DEFINE_string(protocol, "tcp", "Protocol"); DEFINE_int32(port, 50052, "Real Client service port"); DEFINE_string(global_segment_size, "4 GB", "Size of global segment"); +DEFINE_string(local_buffer_size, "0", "Size of local buffer (e.g., 16MB, 1GB)"); DEFINE_int32(threads, 1, "Number of threads for client service"); DEFINE_string(tenant_id, "default", "Tenant identifier"); DEFINE_bool(enable_offload, false, "Enable offload availability"); @@ -103,6 +104,7 @@ int main(int argc, char *argv[]) { } size_t global_segment_size = string_to_byte_size(FLAGS_global_segment_size); + size_t local_buffer_size = string_to_byte_size(FLAGS_local_buffer_size); #ifdef USE_ASCEND_DIRECT // just set to true, does not affect GPU process. globalConfig().ascend_agent_mode = true; @@ -110,10 +112,11 @@ int main(int argc, char *argv[]) { auto client_inst = RealClient::create(); auto res = client_inst->setup_internal( - FLAGS_host, FLAGS_metadata_server, global_segment_size, 0, - FLAGS_protocol, FLAGS_device_names, FLAGS_master_server_address, - nullptr, "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock", - FLAGS_port, FLAGS_enable_offload, FLAGS_start_offload_rpc_server, "", + FLAGS_host, FLAGS_metadata_server, global_segment_size, + local_buffer_size, FLAGS_protocol, FLAGS_device_names, + FLAGS_master_server_address, nullptr, + "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock", FLAGS_port, + FLAGS_enable_offload, FLAGS_start_offload_rpc_server, "", FLAGS_tenant_id); if (!res) { LOG(FATAL) << "Failed to setup client: " << toString(res.error()); From 84b2d91c1da1a81db951d7882aa8bf80b3e39af7 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Tue, 7 Jul 2026 18:37:47 +0800 Subject: [PATCH 29/57] [TENT] Opt-in earliest-deadline-first dispatch in admission queue (RFC #2519 step 2) (#2763) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Opt-in earliest-deadline-first dispatch in admission queue (RFC #2519 step 2) Step 1 (#2618) added Request.deadline_ns + post-hoc MLU observability. This adds the first policy step alogfans green-lit: an opt-in EDF ordering in LocalTransferAdmissionQueue::pickForDispatch. - QueueLimits gains `deadline_aware` (default false). - When false: pickForDispatch keeps strict FIFO — behavior unchanged. - When true: the dispatch queue is stably reordered earliest-deadline-first before selection; owners without a deadline (deadline_ns == 0) keep FIFO order behind all deadlined owners. - Selection still respects the existing owner/byte capacity limits; this only reorders *which* queued owner is picked next, it does not admit or reject. No codec / local-decode / bandwidth-prediction here — this is the ordering layer only, strictly additive and gated behind the opt-in flag. Motivation from measurement (H20 / CX-7 RoCEv2, TENT backend): sweeping the deadline shows a transition band (~200us on this HW) where feasible and missed transfers coexist (mean MLU 1.31, ~13% feasible) — exactly where EDF ordering has leverage. Below/above that band ordering buys nothing. Misses there are driven by concurrency queueing, which is what this reordering targets. Adds 3 unit tests (EDF order, undeadlined-last, and FIFO-unchanged default); full admission_queue_test suite passes (16/16). Co-Authored-By: Claude Opus 4.8 * [TENT] EDF: order fifo_ at admission instead of re-sorting every dispatch Addresses review on #2763: pickForDispatch used to std::stable_sort the whole fifo_ (up to max_outstanding_owners, default 1024) on every call, and it is called on every complete/poll/submit — re-sorting an already-ordered queue repeatedly, even when nothing new was admitted, and even when a byte limit lets it consume only a few entries. Instead keep fifo_ EDF-ordered as owners arrive: when deadline_aware, tryAdmit inserts each owner at its earliest-deadline-first position (upper_bound, so same-deadline owners keep FIFO tie-break — identical ordering to the old stable sort). pickForDispatch then just consumes from the front, dropping the sort entirely. Hot dispatch path goes from O(N log N) to O(picked); the cost moves to one O(N) ordered insert per admit. Default (deadline_aware == false) is unchanged plain FIFO push_back. Adds a test admitting out-of-order deadlines across separate tryAdmit calls to cover the ordered-insert path; full admission_queue_test passes. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 --- .../include/tent/runtime/admission_queue.h | 7 ++ .../tent/src/runtime/admission_queue.cpp | 37 +++++- .../tent/tests/admission_queue_test.cpp | 109 ++++++++++++++++++ 3 files changed, 152 insertions(+), 1 deletion(-) diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index e52f70602a..d5e970b15c 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -42,6 +42,13 @@ struct QueueLimits { size_t max_outstanding_bytes{0}; size_t staging_owner_reserve{0}; size_t staging_byte_reserve{0}; + // Opt-in deadline-aware dispatch (RFC #2519 step 2). When false (default), + // pickForDispatch keeps strict FIFO order — unchanged behavior. When true, + // owners carrying a deadline (request.deadline_ns != 0) are dispatched + // earliest-deadline-first; owners without a deadline keep FIFO order behind + // them. This only reorders selection within the existing capacity limits; + // it does not admit/reject or otherwise change what gets dispatched. + bool deadline_aware{false}; }; struct QueueOwnerInput { diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 3d30747a77..e236e54a5c 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -14,6 +14,7 @@ #include "tent/runtime/admission_queue.h" +#include #include #include @@ -23,6 +24,13 @@ namespace { using PublicTaskKey = std::pair; +// Sort key for EDF: owners without a deadline (0) sort after all deadlined +// owners, so they never jump ahead of a real deadline. +inline uint64_t deadlineKey(uint64_t deadline_ns) { + return deadline_ns == 0 ? std::numeric_limits::max() + : deadline_ns; +} + bool isSupportedTerminalStatus(TransferStatusEnum status) { return status == TransferStatusEnum::COMPLETED || status == TransferStatusEnum::INVALID || @@ -174,7 +182,27 @@ Status LocalTransferAdmissionQueue::tryAdmit( for (const auto derived_task_id : owner_input.derived_task_ids) { public_to_owner_[{submit.batch_token, derived_task_id}] = owner_id; } - fifo_.push_back(owner_id); + // RFC #2519 step 2: keep fifo_ ordered on admission so pickForDispatch + // never has to re-sort. Default (deadline_aware == false) appends in + // strict FIFO. When deadline-aware, insert at the earliest-deadline- + // first position; upper_bound places a new owner *after* existing + // owners with the same deadline, preserving FIFO order among ties. + if (limits_.deadline_aware) { + const uint64_t key = deadlineKey(owner.request.deadline_ns); + auto pos = std::upper_bound( + fifo_.begin(), fifo_.end(), key, + [this](uint64_t k, QueueOwnerId id) { + auto it = owners_.find(id); + uint64_t d = + (it == owners_.end()) + ? std::numeric_limits::max() + : deadlineKey(it->second.request.deadline_ns); + return k < d; + }); + fifo_.insert(pos, owner_id); + } else { + fifo_.push_back(owner_id); + } admitted_owner_ids.push_back(owner_id); } @@ -190,6 +218,13 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( std::vector picked; if (max_owners == 0 || max_bytes == 0) return picked; + // RFC #2519 step 2 (opt-in): earliest-deadline-first dispatch. When + // deadline_aware, fifo_ is kept EDF-ordered at admission time (see + // tryAdmit's ordered insert), so there is nothing to sort here — we just + // consume from the front. This keeps the hot dispatch path O(picked) + // instead of re-sorting the whole queue (up to max_outstanding_owners) on + // every call. Default (deadline_aware == false) is plain FIFO. + size_t used_owners = 0; size_t used_bytes = 0; while (!fifo_.empty() && used_owners < max_owners) { diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index b26ff889f9..f22edce8c3 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -388,6 +388,115 @@ TEST(AdmissionQueueTest, AllowsBatchTokenReuseAfterRetire) { EXPECT_EQ(resolved_owner, 2u); } +// --- RFC #2519 step 2: opt-in deadline-aware (EDF) dispatch --------------- + +QueueOwnerInput makeOwnerWithDeadline(size_t public_task_id, size_t length, + uint64_t deadline_ns) { + QueueOwnerInput owner = makeOwner(public_task_id, length); + owner.request.deadline_ns = deadline_ns; + return owner; +} + +TEST(AdmissionQueueTest, DeadlineAwareDispatchesEarliestDeadlineFirst) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + // Admitted in FIFO order 1,2,3 but with deadlines 300,100,200. + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 300), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 200)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + ASSERT_EQ(admitted_ids.size(), 3u); // owner ids 1,2,3 + + auto picked = queue.pickForDispatch(3, 4096); + // EDF order: owner 2 (dl 100) < owner 3 (dl 200) < owner 1 (dl 300). + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeadlineAwareKeepsUndeadlinedOwnersLast) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + // owner 1: no deadline (0); owner 2: deadline 100; owner 3: no deadline. + auto status = queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 0), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 0)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(3, 4096); + // Deadlined owner 2 first; undeadlined 1,3 keep FIFO order behind it. + const std::vector expected{2, 1, 3}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeadlineUnawareKeepsStrictFifo) { + // Default (deadline_aware == false): FIFO regardless of deadlines. + LocalTransferAdmissionQueue queue({4, 4096, 0, 0}); + std::vector admitted_ids; + + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 300), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 200)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(3, 4096); + const std::vector expected{1, 2, + 3}; // FIFO, deadlines ignored + EXPECT_EQ(picked, expected); +} + +// fifo_ is kept EDF-ordered at admission time, so owners admitted across +// *separate* tryAdmit calls (out of deadline order) must still dispatch EDF — +// this exercises the ordered-insert path, not just a single sorted batch. +TEST(AdmissionQueueTest, DeadlineAwareOrdersAcrossSeparateAdmits) { + QueueLimits limits{8, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector ids; + + // Admit one at a time, deadlines arriving out of order: 300, 100, 200, 0. + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 300)}), + ids) + .code(), + Status::Code::kOk); // owner 1 + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 100)}), + ids) + .code(), + Status::Code::kOk); // owner 2 + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 200)}), + ids) + .code(), + Status::Code::kOk); // owner 3 + ASSERT_EQ( + queue.tryAdmit(makeSubmit(4, 1, {makeOwnerWithDeadline(0, 16, 0)}), ids) + .code(), + Status::Code::kOk); // owner 4 (no deadline → last) + + auto picked = queue.pickForDispatch(8, 4096); + // EDF: 100(owner2) < 200(owner3) < 300(owner1) < no-deadline(owner4). + const std::vector expected{2, 3, 1, 4}; + EXPECT_EQ(picked, expected); +} + } // namespace } // namespace tent } // namespace mooncake From 0d7f0a02f773389b3e1da988e0e7cb520ff65dc9 Mon Sep 17 00:00:00 2001 From: Feng Ren Date: Wed, 8 Jul 2026 10:25:37 +0800 Subject: [PATCH 30/57] [Bench] Enable replay speedup and multi-threading in SSD Benchmarking (#2780) * [Bench] Enable replay speedup and multi-threading * Correct multithreading * Update benchmarks/storage_benchmark_v1/benchmark.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Change multi-thread benchmarking * Add notes --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../storage_benchmark/storage_benchmark.py | 956 ------------------ benchmarks/storage_benchmark_v1/benchmark.py | 447 ++++++-- benchmarks/storage_benchmark_v1/doc/README.md | 73 +- .../storage_benchmark_v1/storage/disk.py | 13 +- 4 files changed, 435 insertions(+), 1054 deletions(-) delete mode 100644 benchmarks/storage_benchmark/storage_benchmark.py diff --git a/benchmarks/storage_benchmark/storage_benchmark.py b/benchmarks/storage_benchmark/storage_benchmark.py deleted file mode 100644 index 689bc98e99..0000000000 --- a/benchmarks/storage_benchmark/storage_benchmark.py +++ /dev/null @@ -1,956 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 - -""" -Mooncake KVCache Storage Benchmark Tool -""" - -import argparse -import json -import time -import os -import statistics -import random -import errno -from pathlib import Path -from typing import Dict, List, Optional -from dataclasses import dataclass - -# ============================================================================ -# Constants -# ============================================================================ - -BLOCK_SIZE_TOKENS = 512 # Number of tokens per block -DEFAULT_BYTES_PER_TOKEN = 2048 # 7B model FP16 (2KB per token) -BLOCK_SIZE_BYTES = BLOCK_SIZE_TOKENS * DEFAULT_BYTES_PER_TOKEN # 1MB per block -MIN_LATENCY_MS = 0.001 # Minimum latency in milliseconds (1 microsecond) - -# Model KVCache sizes (bytes per token, based on LMCache calculator) -# Source: https://lmcache.ai/kv_cache_calculator.html -MODEL_BYTES_PER_TOKEN = { - "llama-3.1-405b": 327680, - "qwen3-32b": 81920, - "deepseek-v3": 1748992, - "glm-4.6": 157013, - "default": DEFAULT_BYTES_PER_TOKEN, -} - -# ============================================================================ -# Data Structures -# ============================================================================ - -@dataclass -class KVCacheRequest: - """KVCache request - - Attributes: - timestamp: Request timestamp in milliseconds - hash_ids: List of block IDs (each ID corresponds to a 512-token block) - input_length: Input token count - output_length: Output token count - """ - timestamp: float - hash_ids: List[int] - input_length: int - output_length: int - -# ============================================================================ -# Storage Layer: Offset Allocator -# ============================================================================ - -class OffsetAllocatorStorage: - """High-performance block storage based on Offset Allocator - - Architecture: - ----------- - 1. Single large file stores all blocks (avoids file explosion) - 2. Uses offset to manage file space (similar to Mooncake's OffsetAllocator) - 3. hash_id -> offset mapping stored in memory (fast lookup) - - Block Organization: - ----------- - Each block corresponds to 512 tokens, fixed size 1MB: - - hash_id[0] -> block_0 (tokens [0...511]) -> offset 0 - - hash_id[1] -> block_1 (tokens [512...1023]) -> offset 1 - - hash_id[i] -> block_i (tokens [i*512...(i+1)*512-1]) -> offset i - - Performance Advantages: - ----------- - - Only one file, no file explosion - - Offset reuse, reduces memory allocation - - pread/pwrite, thread-safe, no seek needed - - Keep fd open, reduces open/close overhead - - Metadata in memory, O(1) lookup - - Attributes: - storage_dir: Storage directory path - block_size_bytes: Block size in bytes - max_blocks: Maximum number of blocks - hash_id_to_offset: hash_id -> offset mapping - free_offsets: List of reusable offsets - next_offset: Next allocatable offset - """ - - def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_blocks: int = 100000, block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100): - """Initialize Offset Allocator storage - - Args: - storage_dir: Storage directory path - bytes_per_token: Bytes per token - max_blocks: Maximum number of blocks (determines file size) - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - """ - self.storage_dir = Path(storage_dir) - self.bytes_per_token = bytes_per_token - self.block_size_tokens = block_size_tokens - self.block_size_bytes = self.block_size_tokens * self.bytes_per_token - self.max_blocks = max_blocks - - # Fsync configuration - self.fsync_mode = fsync_mode - self.fsync_batch_size = fsync_batch_size - self.pending_sync_count = 0 - - # Create storage directory - self.storage_dir.mkdir(parents=True, exist_ok=True) - - # Single large file - self.storage_file = self.storage_dir / "kvcache_storage.bin" - self.file_size = self.max_blocks * self.block_size_bytes - - # Initialize storage file - if not self.storage_file.exists(): - self._init_storage_file() - - # hash_id -> offset mapping (metadata, in memory) - self.hash_id_to_offset: Dict[int, int] = {} - - # Offset allocator (free list) - self.free_offsets: List[int] = [] - self.next_offset = 0 - - # File descriptor (keep open, avoid repeated open/close) - self.fd = None - - # Pre-allocated data buffer with pattern to avoid SSD compression artifacts - # Using a repeating pattern that looks like realistic data (not all zeros) - # Pattern: 64-byte repeated sequence mixed with some variation - pattern = bytes([(i & 0xFF) for i in range(256)]) # 0-255 byte pattern - pattern_repeats = (self.block_size_bytes // len(pattern)) + 1 - self._data_buffer = (pattern * pattern_repeats)[:self.block_size_bytes] - - # Statistics - self.stats = { - 'read_count': 0, - 'write_count': 0, - 'read_bytes': 0, - 'write_bytes': 0, - 'read_latencies_ms': [], - 'write_latencies_ms': [], - 'sync_count': 0, # Number of fsync operations performed - } - - # ======================================================================== - # Internal Methods - # ======================================================================== - - def _init_storage_file(self): - """Initialize storage file (pre-allocate space) - - Create sparse file to avoid actual disk space usage until data is written - """ - with open(self.storage_file, 'wb') as f: - f.seek(self.file_size - 1) - f.write(b'\0') - f.flush() - os.fsync(f.fileno()) - - def _get_fd(self): - """Get file descriptor (lazy open) - - Returns: - int: File descriptor - """ - if self.fd is None: - # Use O_RDWR | O_CREAT, no O_DIRECT (Python compatibility) - self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT) - return self.fd - - def _allocate_offset(self) -> int: - """Allocate a new offset - - Prioritize reusing freed offsets, otherwise allocate new offset - - Returns: - int: Allocated offset - """ - if self.free_offsets: - return self.free_offsets.pop() - offset = self.next_offset - self.next_offset += 1 - return offset - - def _free_offset(self, offset: int): - """Free offset for reuse - - Args: - offset: Offset to free - """ - self.free_offsets.append(offset) - - # ======================================================================== - # Public Interface - # ======================================================================== - - def block_exists(self, hash_id: int) -> bool: - """Check if block exists - - Args: - hash_id: Unique block identifier - - Returns: - bool: Whether block exists - """ - return hash_id in self.hash_id_to_offset - - def read_block(self, hash_id: int) -> float: - """Read block using pread - - Args: - hash_id: Unique block identifier - - Returns: - float: Read latency in milliseconds, or 0 if block doesn't exist - """ - if hash_id not in self.hash_id_to_offset: - return 0.0 # Block doesn't exist, no latency to measure - - offset = self.hash_id_to_offset[hash_id] - file_offset = offset * self.block_size_bytes - - start = time.perf_counter() - - try: - fd = self._get_fd() - data = os.pread(fd, self.block_size_bytes, file_offset) - latency_ms = (time.perf_counter() - start) * 1000.0 - - self.stats['read_count'] += 1 - self.stats['read_bytes'] += len(data) - self.stats['read_latencies_ms'].append(latency_ms) - return latency_ms - except OSError as e: - print(f"Error reading block {hash_id} at offset {file_offset}: {e}") - return 0.0 # Error case, don't pollute stats - - def write_block(self, hash_id: int) -> float: - """Write block using pwrite - - Args: - hash_id: Unique block identifier - - Returns: - float: Write latency in milliseconds - """ - # Allocate offset - offset = self._allocate_offset() - file_offset = offset * self.block_size_bytes - - # Use pre-allocated buffer (much faster than os.urandom) - data = self._data_buffer - - start = time.perf_counter() - - try: - fd = self._get_fd() - written = os.pwrite(fd, data, file_offset) - - write_done = time.perf_counter() - - # Conditional fsync based on mode - if self.fsync_mode == 'always': - # Include fsync in latency measurement - os.fsync(fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - latency_ms = (time.perf_counter() - start) * 1000.0 - # Evict from page cache AFTER fsync to ensure reads measure actual SSD performance - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - elif self.fsync_mode == 'batch': - # For batch mode, only measure write time (fsync is deferred) - self.pending_sync_count += 1 - if self.pending_sync_count >= self.fsync_batch_size: - os.fsync(fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - latency_ms = (write_done - start) * 1000.0 # Only write time - # Evict from page cache after each write - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - elif self.fsync_mode == 'none': - latency_ms = (write_done - start) * 1000.0 - # Evict from page cache even when not syncing - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - else: # 'end' mode - latency_ms = (write_done - start) * 1000.0 - # Evict from page cache (fsync will happen at the end) - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - - # Update mapping - self.hash_id_to_offset[hash_id] = offset - - self.stats['write_count'] += 1 - self.stats['write_bytes'] += written - self.stats['write_latencies_ms'].append(latency_ms) - return latency_ms - except OSError as e: - if e.errno == errno.ENOSPC: - print(f"Error: Disk full when writing block {hash_id} at offset {file_offset}") - else: - print(f"Error writing block {hash_id} at offset {file_offset}: {e}") - return 0.0 # Error case, don't pollute stats - - def __enter__(self): - """Context manager entry""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - ensures cleanup""" - # Perform final fsync before closing for 'end' and 'batch' modes - self._finalize_sync() - self.close(force_sync=False) # Already synced above - return False - - def _finalize_sync(self): - """Perform final fsync before closing (for 'end' mode and pending batch writes)""" - if self.fd is not None: - if self.fsync_mode == 'end': - try: - os.fsync(self.fd) - self.stats['sync_count'] += 1 - except OSError: - pass - elif self.fsync_mode == 'batch' and self.pending_sync_count > 0: - # Flush remaining pending writes - try: - os.fsync(self.fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - except OSError: - pass - - def close(self, force_sync: bool = True): - """Close file - - Args: - force_sync: Whether to force fsync before closing - """ - # For backward compatibility with non-context-manager usage - if force_sync: - self._finalize_sync() - - if self.fd is not None: - os.close(self.fd) - self.fd = None - - def get_stats(self) -> Dict: - """Get statistics - - Returns: - Dict: Dictionary containing read/write statistics - """ - def calc_stats(latencies): - """Calculate latency statistics""" - if not latencies: - return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - return { - 'avg_ms': statistics.mean(latencies), - **calc_percentiles(latencies), - } - - return { - 'read': { - 'count': self.stats['read_count'], - 'mb': self.stats['read_bytes'] / 1024 / 1024, - **calc_stats(self.stats['read_latencies_ms']) - }, - 'write': { - 'count': self.stats['write_count'], - 'mb': self.stats['write_bytes'] / 1024 / 1024, - **calc_stats(self.stats['write_latencies_ms']) - }, - 'sync_count': self.stats['sync_count'], - 'total_blocks': len(self.hash_id_to_offset), - 'free_blocks': len(self.free_offsets), - } - - -# ============================================================================ -# Benchmark Layer -# ============================================================================ - -class StorageBenchmark: - """KVCache storage benchmark - - Based on Mooncake OffsetAllocator + vLLM PagedAttention implementation: - - Example: - ----- - Request A: [1, 2, 4] - -> hash_id 1 -> not exist, write block_1 (offset=0, 1MB) - -> hash_id 2 -> not exist, write block_2 (offset=1, 1MB) - -> hash_id 4 -> not exist, write block_4 (offset=2, 1MB) - - Request B: [1, 2, 4, 6] - -> hash_id 1 -> exists, read block_1 (offset=0) ✓ prefix reuse - -> hash_id 2 -> exists, read block_2 (offset=1) ✓ prefix reuse - -> hash_id 4 -> exists, read block_4 (offset=2) ✓ prefix reuse - -> hash_id 6 -> not exist, write block_6 (offset=3, 1MB) - - Performance Advantages: - --------- - - Single file operation, no file explosion - - Offset reuse, reduces memory allocation - - pread/pwrite, thread-safe - """ - - def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_blocks: int = 100000, block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100): - """Initialize benchmark - - Args: - storage_dir: Storage directory - bytes_per_token: Bytes per token - max_blocks: Maximum number of blocks - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - """ - self.storage = OffsetAllocatorStorage( - storage_dir, bytes_per_token, max_blocks, - block_size_tokens, fsync_mode, fsync_batch_size - ) - self.bytes_per_token = bytes_per_token - self.block_size_tokens = block_size_tokens - - # Statistics - self.stats = { - 'total_requests': 0, - 'total_blocks': 0, - 'read_blocks': 0, - 'write_blocks': 0, - 'prefix_hit_blocks': 0, # Number of prefix hit blocks - 'request_latencies_ms': [], - } - - def process_request(self, req: KVCacheRequest) -> float: - """Process a KVCache request - - Based on vLLM's prefix caching mechanism: - - Each hash_id corresponds to an independent block - - Prefix reuse achieved through hash_id matching - - Args: - req: KVCache request - - Returns: - float: Request latency in milliseconds - """ - self.stats['total_requests'] += 1 - self.stats['total_blocks'] += len(req.hash_ids) - - start_time = time.perf_counter() - total_latency = 0.0 - - # Process each hash_id (in order) - for hash_id in req.hash_ids: - if self.storage.block_exists(hash_id): - # Block exists, read (reuse cached block) - total_latency += self.storage.read_block(hash_id) - self.stats['read_blocks'] += 1 - self.stats['prefix_hit_blocks'] += 1 # Count all cache hits as prefix reuse - else: - # Block doesn't exist, write (new block) - total_latency += self.storage.write_block(hash_id) - self.stats['write_blocks'] += 1 - - latency_ms = total_latency if total_latency > 0 else MIN_LATENCY_MS - self.stats['request_latencies_ms'].append(latency_ms) - - return latency_ms - - def get_stats(self) -> Dict: - """Get statistics - - Returns: - Dict: Statistics dictionary - """ - storage_stats = self.storage.get_stats() - - request_latencies = self.stats['request_latencies_ms'] - - if request_latencies: - latency_stats = { - 'avg_ms': statistics.mean(request_latencies), - **calc_percentiles(request_latencies), - } - else: - latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - - total_blocks = self.stats['total_blocks'] - read_blocks = self.stats['read_blocks'] - write_blocks = self.stats['write_blocks'] - - return { - 'total_requests': self.stats['total_requests'], - 'total_blocks': total_blocks, - 'read_blocks': read_blocks, - 'write_blocks': write_blocks, - 'prefix_hit_blocks': self.stats['prefix_hit_blocks'], - 'block_hit_rate': read_blocks / total_blocks if total_blocks > 0 else 0, - 'write_ratio': write_blocks / total_blocks if total_blocks > 0 else 0, - 'tokens_per_block': self.block_size_tokens, # Configurable block size in tokens - 'latency': latency_stats, - 'storage': storage_stats, - } - - def __enter__(self): - """Context manager entry""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - ensures cleanup""" - self.close() - return False - - def close(self, force_sync: bool = True): - """Close storage - - Args: - force_sync: Whether to force final sync before closing - """ - self.storage.close(force_sync=force_sync) - - -# ============================================================================ -# Utility Functions -# ============================================================================ - -def calc_percentiles(data: List[float]) -> Dict[str, float]: - """Calculate latency percentiles - - Uses linear interpolation for accurate percentile calculation. - This is more accurate than statistics.quantiles() for small datasets. - - Args: - data: List of latency values in milliseconds - - Returns: - Dict containing p50, p95, p99 percentiles - """ - if not data: - return {'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - - # Sort data for percentile calculation - sorted_data = sorted(data) - n = len(sorted_data) - - def get_percentile(p: float) -> float: - """Get percentile using linear interpolation - - Args: - p: Percentile (0-100) - - Returns: - Value at percentile - """ - index = (n - 1) * p / 100 - lower = int(index) - upper = min(lower + 1, n - 1) - - if lower == upper: - return sorted_data[lower] - - # Linear interpolation - weight = index - lower - return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight - - return { - 'p50_ms': get_percentile(50), - 'p95_ms': get_percentile(95), - 'p99_ms': get_percentile(99), - } - - -# ============================================================================ -# Trace Loader -# ============================================================================ - -class TraceLoader: - """Load KVCache trace""" - - def __init__(self, trace_path: str): - """Initialize trace loader - - Args: - trace_path: Trace file path - """ - self.trace_path = trace_path - self.requests = [] - self._load_trace() - - def _load_trace(self): - """Load trace file with error handling""" - line_num = 0 - try: - with open(self.trace_path, 'r') as f: - for line in f: - line_num += 1 - line = line.strip() - if not line: - continue - try: - req = json.loads(line) - # Validate required fields - if not all(k in req for k in ['timestamp', 'hash_ids', 'input_length', 'output_length']): - print(f"Warning: Line {line_num} missing required fields, skipping") - continue - if not isinstance(req['hash_ids'], list): - print(f"Warning: Line {line_num} has invalid hash_ids (not a list), skipping") - continue - self.requests.append(KVCacheRequest( - timestamp=float(req['timestamp']), - hash_ids=req['hash_ids'], - input_length=int(req['input_length']), - output_length=int(req['output_length']) - )) - except (json.JSONDecodeError, ValueError, KeyError) as e: - print(f"Warning: Line {line_num} has invalid format: {e}, skipping") - continue - except FileNotFoundError: - raise FileNotFoundError(f"Trace file not found: {self.trace_path}") - except OSError as e: - raise OSError(f"Error reading trace file {self.trace_path}: {e}") - - def get_requests(self) -> List[KVCacheRequest]: - """Get request list - - Returns: - List[KVCacheRequest]: Request list - """ - return self.requests - - -# ============================================================================ -# Benchmark Runner -# ============================================================================ - -def run_benchmark(trace_path: str, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_requests: Optional[int] = None, max_blocks: int = 100000, - replay_timestamps: bool = False, time_scale: float = 1.0, - block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100) -> Dict: - """Run benchmark - - Args: - trace_path: Trace file path - storage_dir: Storage directory - bytes_per_token: Bytes per token - max_requests: Maximum number of requests (None = all) - max_blocks: Maximum number of blocks - replay_timestamps: Whether to replay timestamps from trace (simulate realistic timing) - time_scale: Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed) - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - - Returns: - Dict: Benchmark results - """ - block_size_bytes = block_size_tokens * bytes_per_token - - print(f"\n{'='*80}") - print(f"Running: {Path(trace_path).name}") - print(f"Architecture: Offset Allocator (Mooncake style)") - print(f"Block size: {block_size_tokens} tokens/block ({block_size_bytes:,} bytes)") - print(f"Storage: Single large file with offset-based block management") - print(f"Bytes per token: {bytes_per_token}") - print(f"Max blocks: {max_blocks}") - print(f"Fsync mode: {fsync_mode}" + (f" (batch_size={fsync_batch_size})" if fsync_mode == 'batch' else '')) - print(f"Timestamp replay: {'Enabled' if replay_timestamps else 'Disabled'}") - if replay_timestamps: - scale_desc = 'real-time' if time_scale == 1.0 else f'{1/time_scale:.1f}x speed' if time_scale < 1.0 else f'{time_scale}x slower' - print(f"Time scale: {time_scale}x ({scale_desc})") - print(f"{'='*80}") - - # Load trace - loader = TraceLoader(trace_path) - requests = loader.get_requests() - - if max_requests: - requests = requests[:max_requests] - - print(f"Loaded {len(requests)} requests") - - # Show timestamp range - if replay_timestamps and requests: - timestamps = [req.timestamp for req in requests] - time_span_ms = max(timestamps) - min(timestamps) - print(f"Timestamp range: {min(timestamps):.1f} - {max(timestamps):.1f} ms (span: {time_span_ms:.1f} ms)") - - # Create benchmark instance with context manager for cleanup - with StorageBenchmark( - storage_dir, bytes_per_token, max_blocks, - block_size_tokens, fsync_mode, fsync_batch_size - ) as benchmark: - - # Run benchmark - start_time = time.perf_counter() - total_io_time = 0.0 # Actual I/O time (excluding sleep) - last_timestamp = None - base_time = time.time() # Use wall time for replay synchronization - - for i, req in enumerate(requests): - # Replay by timestamps - sleep_time = 0.0 - if replay_timestamps and last_timestamp is not None: - # Calculate time interval from previous request - delta_ms = req.timestamp - last_timestamp - sleep_time = delta_ms / 1000.0 / time_scale # Apply time scaling - - if sleep_time > 0: - time.sleep(sleep_time) - - # Process request (measure I/O time) - req_start = time.perf_counter() - benchmark.process_request(req) - req_io_time = time.perf_counter() - req_start - total_io_time += req_io_time - - # Record current request timestamp - last_timestamp = req.timestamp - - # Progress output - if (i + 1) % 100 == 0: - if replay_timestamps: - elapsed_wall_time = time.time() - base_time - simulated_time = (req.timestamp - requests[0].timestamp) / 1000.0 / time_scale - print(f" Processed {i + 1}/{len(requests)}... (wall: {elapsed_wall_time:.1f}s, simulated: {simulated_time:.1f}s, io: {total_io_time:.1f}s)") - else: - print(f" Processed {i + 1}/{len(requests)}...") - - elapsed = time.perf_counter() - start_time - - # Perform final sync to include it in stats - benchmark.storage._finalize_sync() - - # Get statistics (context manager will handle cleanup) - stats = benchmark.get_stats() - - # Calculate actual I/O time (excluding sleep) - io_time = total_io_time if replay_timestamps else elapsed - - return { - 'trace_file': Path(trace_path).name, - 'total_requests': len(requests), - 'simulation_time_s': elapsed, - 'io_time_s': io_time, # Actual I/O time - 'wall_time_s': elapsed, # Wall time (including sleep) - 'requests_per_second': len(requests) / io_time if io_time > 0 else 0, # Based on I/O time - 'timestamp_replay_enabled': replay_timestamps, - 'time_scale': time_scale, - 'bytes_per_token': bytes_per_token, - 'block_size_tokens': block_size_tokens, - 'fsync_mode': fsync_mode, - **stats, - } - - -# ============================================================================ -# Result Output -# ============================================================================ - -def print_results(results: List[Dict]): - """Print benchmark results - - Args: - results: List of benchmark results - """ - for i, r in enumerate(results, 1): - print(f"\n{'='*80}") - print(f" [{i}/{len(results)}] {r['trace_file']}") - print(f"{'='*80}") - - print(f"\n[Performance Overview]") - print(f" Total Requests: {r['total_requests']:,}") - print(f" Queries Per Second (QPS): {r['requests_per_second']:.2f}") - print(f" Cache Hit Rate: {r['block_hit_rate']:.2%}") - print(f" Write Ratio: {r['write_ratio']:.2%}") - print(f" Total Blocks: {r['total_blocks']:,}") - print(f" Read Blocks: {r['read_blocks']:,}") - print(f" Write Blocks: {r['write_blocks']:,}") - print(f" Prefix Hits: {r['prefix_hit_blocks']:,}") - - print(f"\n[Latency Analysis]") - req_lat = r['latency'] - print(f" Request Latency (End-to-End): Avg={req_lat['avg_ms']:.2f}ms, P50={req_lat['p50_ms']:.2f}ms, P95={req_lat['p95_ms']:.2f}ms, P99={req_lat['p99_ms']:.2f}ms") - read_lat = r['storage']['read'] - write_lat = r['storage']['write'] - print(f" Single I/O Operation (Per Block):") - print(f" Read: Avg={read_lat.get('avg_ms', 0):.3f}ms, P50={read_lat.get('p50_ms', 0):.3f}ms, P95={read_lat.get('p95_ms', 0):.3f}ms, P99={read_lat.get('p99_ms', 0):.3f}ms") - print(f" Write: Avg={write_lat.get('avg_ms', 0):.3f}ms, P50={write_lat.get('p50_ms', 0):.3f}ms, P95={write_lat.get('p95_ms', 0):.3f}ms, P99={write_lat.get('p99_ms', 0):.3f}ms") - - print(f"\n[I/O & Bandwidth]") - print(f" Total Read I/O: {r['storage']['read']['mb']:>10.1f} MB ({r['storage']['read']['count']:,} ops)") - print(f" Total Write I/O: {r['storage']['write']['mb']:>10.1f} MB ({r['storage']['write']['count']:,} ops)") - io_time = r['io_time_s'] - bandwidth = (r['storage']['read']['mb'] + r['storage']['write']['mb']) / io_time - print(f" Effective Bandwidth: {bandwidth:>10.1f} MB/s") - - print(f"\n[Storage Details]") - print(f" Blocks in Use: {r['storage']['total_blocks']:>10,}") - print(f" Free Blocks: {r['storage']['free_blocks']:>10,}") - print(f" Tokens per Block: {r['tokens_per_block']:>10,}") - print(f" Block Size: {r['tokens_per_block'] * r.get('bytes_per_token', 2048) / 1024 / 1024:>10.2f} MB") - if 'sync_count' in r['storage']: - print(f" Fsync Operations: {r['storage']['sync_count']:>10,}") - - print(f"\n[Execution Time]") - if r.get('timestamp_replay_enabled'): - print(f" Wall Time (Total): {r['wall_time_s']:>10.2f} s") - print(f" I/O Time (Actual): {r['io_time_s']:>10.2f} s") - print(f" Sleep Time (Replay): {r['wall_time_s'] - r['io_time_s']:>10.2f} s") - else: - print(f" Total Execution Time: {r['wall_time_s']:>10.2f} s") - - print(f"\n{'='*80}\n") - - -# ============================================================================ -# Main Program -# ============================================================================ - -def main(): - """Main entry point""" - parser = argparse.ArgumentParser( - description='Mooncake KVCache Storage Benchmark', - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Examples: - # Quick test (100 requests) - python storage_benchmark.py --scenario=toolagent --max-requests=100 - - # Test with large model preset (Llama-3.1-405B) - python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100 - - # Test with Deepseek V3 (extra large model) - python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100 - - # Realistic replay (with timestamps, 10x speed) - python storage_benchmark.py --scenario=toolagent --max-requests=1000 \\ - --replay-timestamps --time-scale=0.1 - - # All scenarios with custom bytes_per_token - python storage_benchmark.py --scenario=all --bytes-per-token=512 - - # Test with different block sizes and fsync modes - python storage_benchmark.py --scenario=toolagent --block-size-tokens=256 --fsync-mode=always - - # Test with custom fsync batch size - python storage_benchmark.py --scenario=toolagent --fsync-mode=batch --fsync-batch-size=50 - -Performance Tuning: - --fsync-mode=batch (default): Balance between performance and safety - --fsync-mode=always: Safest but slowest, measures full persistence cost - --fsync-mode=end: Fastest, only measures write I/O (not persistence) - --fsync-mode=none: Testing only, no durability guarantees - -Available model presets: - llama-3.1-405b, qwen3-32b, deepseek-v3, glm-4.6, default - -For more information: tools/STORAGE_BENCHMARK_README.md - """ - ) - - parser.add_argument('--trace-dir', type=str, default='../../FAST25-release/traces', - help='Trace files directory') - parser.add_argument('--scenario', type=str, choices=['conversation', 'synthetic', 'toolagent', 'all'], - default='toolagent', help='Test scenario') - parser.add_argument('--storage-dir', type=str, default='/tmp/mooncake_bench', - help='Storage directory') - parser.add_argument('--model', type=str, choices=list(MODEL_BYTES_PER_TOKEN.keys()), - default='default', - help=f'Model preset (overrides --bytes-per-token). Available: {", ".join(MODEL_BYTES_PER_TOKEN.keys())}') - parser.add_argument('--bytes-per-token', type=int, default=DEFAULT_BYTES_PER_TOKEN, - help='Bytes per token (default %d, overridden by --model if specified)' % DEFAULT_BYTES_PER_TOKEN) - parser.add_argument('--max-requests', type=int, default=None, - help='Maximum number of requests (default: unlimited)') - parser.add_argument('--max-blocks', type=int, default=100000, - help='Maximum number of blocks in storage file (determines file size)') - parser.add_argument('--replay-timestamps', action='store_true', - help='Enable timestamp replay (simulate realistic request timing)') - parser.add_argument('--time-scale', type=float, default=1.0, - help='Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)') - parser.add_argument('--block-size-tokens', type=int, default=512, - help='Number of tokens per block (default: 512)') - parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'], - default='batch', - help='When to fsync: batch=every N writes (default), always=after each write, end=only at close, none=never') - parser.add_argument('--fsync-batch-size', type=int, default=100, - help='Number of writes between fsync in batch mode (default: 100)') - - args = parser.parse_args() - - # Print benchmark header - print(f"\n{'='*80}") - print(f"{'Mooncake KVCache Storage Benchmark':^80}") - print(f"{'='*80}") - - # Determine bytes_per_token (model preset takes precedence) - bytes_per_token = MODEL_BYTES_PER_TOKEN.get(args.model, args.bytes_per_token) - if args.model != 'default': - print(f"Using model preset: {args.model} ({bytes_per_token} bytes/token, ~{bytes_per_token/1024:.1f} KB/token)") - else: - print(f"Using custom bytes_per_token: {bytes_per_token}") - - # Determine test scenarios - scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario] - trace_files = { - 'conversation': 'conversation_trace.jsonl', - 'synthetic': 'synthetic_trace.jsonl', - 'toolagent': 'toolagent_trace.jsonl' - } - - # Run benchmarks - results = [] - - for scenario in scenarios: - trace_path = Path(args.trace_dir) / trace_files[scenario] - if trace_path.exists(): - result = run_benchmark( - str(trace_path), - str(Path(args.storage_dir) / scenario), - bytes_per_token, - args.max_requests, - args.max_blocks, - args.replay_timestamps, - args.time_scale, - args.block_size_tokens, - args.fsync_mode, - args.fsync_batch_size - ) - results.append(result) - else: - print(f"Warning: Trace file not found: {trace_path}") - - # Print results - if results: - print_results(results) - - -if __name__ == '__main__': - main() diff --git a/benchmarks/storage_benchmark_v1/benchmark.py b/benchmarks/storage_benchmark_v1/benchmark.py index 3f5ba5ed38..68557d0aab 100644 --- a/benchmarks/storage_benchmark_v1/benchmark.py +++ b/benchmarks/storage_benchmark_v1/benchmark.py @@ -11,6 +11,8 @@ import time import statistics import signal +from contextlib import ExitStack +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from pathlib import Path from typing import Iterator, List, Dict, Any @@ -103,7 +105,8 @@ def __init__(self, storage_dir: str, model_config: dict, 'read_pages': 0, 'write_pages': 0, 'page_hits': 0, - 'request_latencies_ms': [], + 'request_io_latencies_ms': [], + 'request_wall_latencies_ms': [], } def process_request(self, req: KVCacheRequest) -> float: @@ -118,13 +121,14 @@ def process_request(self, req: KVCacheRequest) -> float: self.stats['total_requests'] += 1 self.stats['total_tokens'] += req.input_length + req.output_length - total_latency = 0.0 + request_start = time.perf_counter() + io_latency_ms = 0.0 # Process each access requirement from layout for access in self.layout.get_operations(req): if self.storage.exists(access.page_id): # Page exists, perform READ - total_latency += self.storage.read( + io_latency_ms += self.storage.read( access.page_id, offset_in_page=access.offset_in_page, length=access.length @@ -133,39 +137,23 @@ def process_request(self, req: KVCacheRequest) -> float: self.stats['page_hits'] += 1 else: # Page doesn't exist, perform WRITE - total_latency += self.storage.write( + io_latency_ms += self.storage.write( access.page_id, offset_in_page=access.offset_in_page, length=access.length ) self.stats['write_pages'] += 1 - latency_ms = total_latency if total_latency > 0 else 0.0 - if latency_ms > 0: - self.stats['request_latencies_ms'].append(latency_ms) - return latency_ms + wall_latency_ms = (time.perf_counter() - request_start) * 1000.0 + self.stats['request_io_latencies_ms'].append(io_latency_ms) + self.stats['request_wall_latencies_ms'].append(wall_latency_ms) + return io_latency_ms def get_stats(self) -> Dict: """Get statistics""" storage_stats = self.storage.get_stats() - request_latencies = self.stats['request_latencies_ms'] - - if request_latencies: - sorted_latencies = sorted(request_latencies) - n = len(sorted_latencies) - - def get_percentile(p: float) -> float: - idx = int(n * p) - return sorted_latencies[idx] if idx < n else sorted_latencies[-1] - - latency_stats = { - 'avg_ms': statistics.mean(request_latencies), - 'p50_ms': sorted_latencies[n // 2], - 'p95_ms': get_percentile(0.95), - 'p99_ms': get_percentile(0.99), - } - else: - latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} + request_io_latencies = self.stats['request_io_latencies_ms'] + request_wall_latencies = self.stats['request_wall_latencies_ms'] total_pages = self.stats['read_pages'] + self.stats['write_pages'] @@ -178,7 +166,8 @@ def get_percentile(p: float) -> float: 'page_hits': self.stats['page_hits'], 'page_hit_rate': self.stats['read_pages'] / total_pages if total_pages > 0 else 0, 'write_ratio': self.stats['write_pages'] / total_pages if total_pages > 0 else 0, - 'latency': latency_stats, + 'request_io_latency': latency_stats(request_io_latencies), + 'request_wall_latency': latency_stats(request_wall_latencies), 'storage': storage_stats, } @@ -205,10 +194,228 @@ def get_max_page_id(requests: List[KVCacheRequest]) -> int: return max_id +def parse_csv_floats(value: str) -> List[float]: + return [float(item.strip()) for item in value.split(',') if item.strip()] + + +def wait_for_replay_time(req: KVCacheRequest, base_timestamp: float, + start_time: float, replay_scale: float): + if replay_scale <= 0 or req.timestamp == 0: + return + target_time = (start_time + + max(0.0, req.timestamp - base_timestamp) / + (1000.0 * replay_scale)) + delay = target_time - time.perf_counter() + if delay > 0: + time.sleep(delay) + + +def latency_stats(values: List[float]) -> Dict[str, float]: + if not values: + return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} + + sorted_values = sorted(values) + + def get_percentile(p: float) -> float: + if len(sorted_values) == 1: + return sorted_values[0] + rank = (len(sorted_values) - 1) * p + lower = int(rank) + upper = min(lower + 1, len(sorted_values) - 1) + weight = rank - lower + return (sorted_values[lower] * (1.0 - weight) + + sorted_values[upper] * weight) + + return { + 'avg_ms': statistics.mean(values), + 'p50_ms': get_percentile(0.50), + 'p95_ms': get_percentile(0.95), + 'p99_ms': get_percentile(0.99), + } + + +def snapshot_thread_stats(benchmark: StorageBenchmark) -> Dict[str, Any]: + storage = benchmark.storage + total_pages = benchmark.stats['read_pages'] + benchmark.stats['write_pages'] + return { + 'total_requests': benchmark.stats['total_requests'], + 'total_tokens': benchmark.stats['total_tokens'], + 'read_pages': benchmark.stats['read_pages'], + 'write_pages': benchmark.stats['write_pages'], + 'page_hits': benchmark.stats['page_hits'], + 'request_io_latencies_ms': list( + benchmark.stats['request_io_latencies_ms'] + ), + 'request_wall_latencies_ms': list( + benchmark.stats['request_wall_latencies_ms'] + ), + 'read_bytes': storage.stats['read_bytes'], + 'write_bytes': storage.stats['write_bytes'], + 'read_time_s': storage.stats['read_time_s'], + 'write_time_s': storage.stats['write_time_s'], + 'read_latencies_ms': list(storage.stats['read_latencies_ms']), + 'write_latencies_ms': list(storage.stats['write_latencies_ms']), + 'sync_count': storage.stats['sync_count'], + 'max_pages': storage.max_pages, + 'written_pages': len(storage._written_pages), + 'total_pages': total_pages, + } + + +def aggregate_thread_stats(thread_stats: List[Dict[str, Any]]) -> Dict: + total_requests = sum(s['total_requests'] for s in thread_stats) + total_tokens = sum(s['total_tokens'] for s in thread_stats) + read_pages = sum(s['read_pages'] for s in thread_stats) + write_pages = sum(s['write_pages'] for s in thread_stats) + page_hits = sum(s['page_hits'] for s in thread_stats) + total_pages = read_pages + write_pages + + request_io_latencies = [] + request_wall_latencies = [] + read_latencies = [] + write_latencies = [] + for stats in thread_stats: + request_io_latencies.extend(stats['request_io_latencies_ms']) + request_wall_latencies.extend(stats['request_wall_latencies_ms']) + read_latencies.extend(stats['read_latencies_ms']) + write_latencies.extend(stats['write_latencies_ms']) + + read_bytes = sum(s['read_bytes'] for s in thread_stats) + write_bytes = sum(s['write_bytes'] for s in thread_stats) + read_time = sum(s['read_time_s'] for s in thread_stats) + write_time = sum(s['write_time_s'] for s in thread_stats) + + return { + 'total_requests': total_requests, + 'total_tokens': total_tokens, + 'total_pages': total_pages, + 'read_pages': read_pages, + 'write_pages': write_pages, + 'page_hits': page_hits, + 'page_hit_rate': read_pages / total_pages if total_pages > 0 else 0, + 'write_ratio': write_pages / total_pages if total_pages > 0 else 0, + 'request_io_latency': latency_stats(request_io_latencies), + 'request_wall_latency': latency_stats(request_wall_latencies), + 'storage': { + 'read': { + 'count': read_pages, + 'mb': read_bytes / 1024 / 1024, + 'time_s': read_time, + **latency_stats(read_latencies), + }, + 'write': { + 'count': write_pages, + 'mb': write_bytes / 1024 / 1024, + 'time_s': write_time, + **latency_stats(write_latencies), + }, + 'sync_count': sum(s['sync_count'] for s in thread_stats), + 'max_pages': sum(s['max_pages'] for s in thread_stats), + 'written_pages': sum(s['written_pages'] for s in thread_stats), + 'page_hits': page_hits, + 'page_misses': write_pages, + }, + } + + +def print_progress(done: int, total: int, start_time: float, + stats: Dict, req: KVCacheRequest = None, + suffix: str = ""): + elapsed = time.perf_counter() - start_time + qps = done / elapsed if elapsed > 0 else 0 + storage = stats.get('storage', {}) + read_stats = storage.get('read', {}) + write_stats = storage.get('write', {}) + read_time = read_stats.get('time_s', 0) + write_time = write_stats.get('time_s', 0) + read_mbps = read_stats.get('mb', 0) / read_time if read_time > 0 else 0 + write_mbps = write_stats.get('mb', 0) / write_time if write_time > 0 else 0 + + if req is None: + req_info = "" + else: + req_info = (f" ids={len(req.hash_ids):3d} " + f"tokens={req.input_length + req.output_length:6d} |") + + print(f" [{done:5d}/{total}]{req_info} QPS={qps:7.2f} | " + f"R={stats['read_pages']:6d} " + f"({read_stats.get('avg_ms', 0):6.2f}ms, {read_mbps:6.1f}MB/s) | " + f"W={stats['write_pages']:6d} " + f"({write_stats.get('avg_ms', 0):6.2f}ms, {write_mbps:6.1f}MB/s)" + f"{suffix}") + + +def should_print_progress(done: int, total: int, progress_interval: int) -> bool: + if done >= total: + return True + return progress_interval > 0 and done % progress_interval == 0 + + +def run_single_thread(benchmark: StorageBenchmark, + requests: List[KVCacheRequest], + replay_scale: float, + progress_interval: int) -> Dict[str, Any]: + start_time = time.perf_counter() + base_timestamp = requests[0].timestamp if requests else 0 + completed = 0 + + for req in requests: + wait_for_replay_time(req, base_timestamp, start_time, replay_scale) + benchmark.process_request(req) + completed += 1 + if should_print_progress(completed, len(requests), progress_interval): + print_progress(completed, len(requests), start_time, + benchmark.get_stats(), req) + + return { + 'completed': completed, + 'elapsed': time.perf_counter() - start_time, + 'stats': benchmark.get_stats(), + } + + +def run_multi_thread(benchmarks: List[StorageBenchmark], + requests: List[KVCacheRequest], + replay_scale: float) -> Dict[str, Any]: + start_time = time.perf_counter() + base_timestamp = requests[0].timestamp if requests else 0 + total_requests = len(requests) * len(benchmarks) + completed = 0 + + def run_worker(thread_id: int): + benchmark = benchmarks[thread_id] + for req in requests: + wait_for_replay_time(req, base_timestamp, start_time, replay_scale) + benchmark.process_request(req) + return snapshot_thread_stats(benchmark) + + thread_stats = [] + with ThreadPoolExecutor(max_workers=len(benchmarks)) as executor: + futures = [ + executor.submit(run_worker, thread_id) + for thread_id in range(len(benchmarks)) + ] + for future in as_completed(futures): + worker_stats = future.result() + thread_stats.append(worker_stats) + completed += worker_stats['total_requests'] + print_progress(completed, total_requests, start_time, + aggregate_thread_stats(thread_stats), + suffix=" | completed worker") + + return { + 'completed': completed, + 'elapsed': time.perf_counter() - start_time, + 'stats': aggregate_thread_stats(thread_stats), + } + + def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, max_requests: int = None, max_pages: int = None, page_size_tokens: int = 512, - fsync_mode: str = 'none', fsync_batch_size: int = 100) -> Dict: + fsync_mode: str = 'none', fsync_batch_size: int = 100, + threads: int = 1, replay_scale: float = 0.0, + progress_interval: int = 100) -> Dict: """Run benchmark Args: @@ -220,6 +427,9 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, page_size_tokens: Tokens per page fsync_mode: When to fsync fsync_batch_size: Number of writes between fsync + threads: Benchmark client worker threads + replay_scale: Timestamp replay multiplier; 0 runs unpaced + progress_interval: Print progress every N requests; 0 disables progress Returns: Benchmark results dictionary @@ -229,6 +439,8 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, print(f"Model: {model_config['name']}") print(f"Layers: {model_config['num_layers']}") print(f"Page size: {page_size_tokens} tokens") + print(f"Threads: {threads}") + print(f"Fast-forward: {replay_scale:g}x" if replay_scale > 0 else "Fast-forward: unpaced") print(f"{'='*80}") # Load trace @@ -266,7 +478,11 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, print(f" Pages needed (trace): {max_pages_needed:,}") print(f" Trace storage size: {trace_size_gb:.2f} GB") print(f" Max pages configured: {max_pages:,}") + if threads > 1: + print(f" Max pages across threads: {max_pages * threads:,}") print(f" Max storage available: {max_size_gb:.2f} GB") + if threads > 1: + print(f" Max storage across threads: {max_size_gb * threads:.2f} GB") if max_pages_needed > max_pages: shortfall = max_pages_needed - max_pages @@ -279,65 +495,69 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, surplus_pct = (surplus / max_pages) * 100 if max_pages > 0 else 0 print(f" ✓ Direct mapping: all {max_pages_needed:,} logical pages uniquely mapped") - # Run benchmark - with StorageBenchmark( - storage_dir=storage_dir, - model_config=model_config, - page_size_tokens=page_size_tokens, - max_pages=max_pages, - fsync_mode=fsync_mode, - fsync_batch_size=fsync_batch_size - ) as benchmark: - start_time = time.perf_counter() - try: - for i, req in enumerate(requests): - benchmark.process_request(req) - # Print progress for each request - elapsed = time.perf_counter() - start_time - qps = (i + 1) / elapsed if elapsed > 0 else 0 - stats = benchmark.get_stats() - storage = stats.get('storage', {}) - read_latency = storage.get('read', {}).get('avg_ms', 0) - write_latency = storage.get('write', {}).get('avg_ms', 0) - read_mb = storage.get('read', {}).get('mb', 0) - write_mb = storage.get('write', {}).get('mb', 0) - read_time = storage.get('read', {}).get('time_s', 0) - write_time = storage.get('write', {}).get('time_s', 0) - read_mbps = read_mb / read_time if read_time > 0 else 0 - write_mbps = write_mb / write_time if write_time > 0 else 0 - print(f" [{i+1:5d}/{len(requests)}] ids={len(req.hash_ids):3d} " - f"tokens={req.input_length+req.output_length:6d} | " - f"QPS={qps:7.2f} | " - f"R={stats['read_pages']:6d} ({read_latency:6.2f}ms, {read_mbps:6.1f}MB/s) | " - f"W={stats['write_pages']:6d} ({write_latency:6.2f}ms, {write_mbps:6.1f}MB/s)") - except KeyboardInterrupt: - print(f"\n\n{'='*80}") - print(f"Interrupted! Showing partial results:") - print(f"{'='*80}") - elapsed = time.perf_counter() - start_time - stats = benchmark.get_stats() - print_results([{ - 'trace_file': Path(trace_path).name, - 'total_requests': i + 1, - 'io_time_s': elapsed, - 'requests_per_second': (i + 1) / elapsed if elapsed > 0 else 0, - 'model': model_config['name'], - 'fsync_mode': fsync_mode, - **stats, - }]) - sys.exit(0) - - elapsed = time.perf_counter() - start_time - stats = benchmark.get_stats() + try: + if threads <= 1: + with StorageBenchmark( + storage_dir=storage_dir, + model_config=model_config, + page_size_tokens=page_size_tokens, + max_pages=max_pages, + fsync_mode=fsync_mode, + fsync_batch_size=fsync_batch_size + ) as benchmark: + result = run_single_thread(benchmark, requests, replay_scale, + progress_interval) + else: + with ExitStack() as stack: + benchmarks = [ + stack.enter_context(StorageBenchmark( + storage_dir=str(Path(storage_dir) / f"thread_{thread_id}"), + model_config=model_config, + page_size_tokens=page_size_tokens, + max_pages=max_pages, + fsync_mode=fsync_mode, + fsync_batch_size=fsync_batch_size + )) + for thread_id in range(threads) + ] + result = run_multi_thread(benchmarks, requests, replay_scale) + except KeyboardInterrupt: + print(f"\n\n{'='*80}") + print(f"Interrupted! Showing partial results:") + print(f"{'='*80}") + result = result if 'result' in locals() else { + 'completed': 0, + 'elapsed': 0, + 'stats': {}, + } + print_results([{ + 'trace_file': Path(trace_path).name, + 'total_requests': result['completed'], + 'io_time_s': result['elapsed'], + 'requests_per_second': ( + result['completed'] / result['elapsed'] + if result['elapsed'] > 0 else 0 + ), + 'model': model_config['name'], + 'fsync_mode': fsync_mode, + 'threads': threads, + 'replay_scale': replay_scale, + **result['stats'], + }]) + sys.exit(0) return { 'trace_file': Path(trace_path).name, - 'total_requests': len(requests), - 'io_time_s': elapsed, - 'requests_per_second': len(requests) / elapsed if elapsed > 0 else 0, + 'total_requests': result['completed'], + 'io_time_s': result['elapsed'], + 'requests_per_second': ( + result['completed'] / result['elapsed'] if result['elapsed'] > 0 else 0 + ), 'model': model_config['name'], 'fsync_mode': fsync_mode, - **stats, + 'threads': threads, + 'replay_scale': replay_scale, + **result['stats'], } @@ -350,6 +570,8 @@ def format_storage_stats(stats: Dict, title: str = "Storage"): storage = stats.get('storage', {}) read_stats = storage.get('read', {}) write_stats = storage.get('write', {}) + request_wall = stats.get('request_wall_latency', {}) + request_io = stats.get('request_io_latency', {}) output = [] output.append(f"\n[{title}]") @@ -357,12 +579,28 @@ def format_storage_stats(stats: Dict, title: str = "Storage"): # General info output.append(f"\n[General]") output.append(f" Model: {stats.get('model', 'N/A')}") + output.append(f" Threads: {stats.get('threads', 1)}") + replay_scale = stats.get('replay_scale', 0) + output.append(f" Fast-forward: {f'{replay_scale:g}x' if replay_scale else 'unpaced'}") output.append(f" Requests: {stats.get('total_requests', 0):,}") output.append(f" Tokens: {stats.get('total_tokens', 0):,}") output.append(f" Total I/O Time: {stats.get('io_time_s', 0):.3f} s") output.append(f" QPS: {stats.get('requests_per_second', 0):.2f}") output.append(f" Hit Rate: {stats.get('page_hit_rate', 0):.2%}") + # Request Stats + output.append(f"\n[Request Wall Latency]") + output.append(f" Avg: {request_wall.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {request_wall.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {request_wall.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {request_wall.get('p99_ms', 0):.3f} ms") + + output.append(f"\n[Request Storage I/O Latency]") + output.append(f" Avg: {request_io.get('avg_ms', 0):.3f} ms") + output.append(f" P50: {request_io.get('p50_ms', 0):.3f} ms") + output.append(f" P95: {request_io.get('p95_ms', 0):.3f} ms") + output.append(f" P99: {request_io.get('p99_ms', 0):.3f} ms") + # Read Stats output.append(f"\n[Read Operations]") output.append(f" Count: {read_stats.get('count', 0):,}") @@ -438,8 +676,18 @@ def main(): default='none', help='When to fsync') parser.add_argument('--fsync-batch-size', type=int, default=100, help='Number of writes between fsync') + parser.add_argument('--threads', type=int, default=1, + help='Number of benchmark client worker threads') + parser.add_argument('--replay-scales', type=str, default='0', + help='Comma-separated trace fast-forward speeds; 0 means unpaced') + parser.add_argument('--progress-interval', type=int, default=100, + help='Print progress every N requests; 0 disables per-request progress') args = parser.parse_args() + if args.threads < 1: + parser.error('--threads must be at least 1') + if args.progress_interval < 0: + parser.error('--progress-interval must be non-negative') print(f"\n{'='*80}") print(f"{'Mooncake KVCache Storage Benchmark':^80}") @@ -447,6 +695,11 @@ def main(): model_config = get_model_config(args.model) print(f"Model: {args.model} ({model_config['num_layers']} layers)") + replay_scales = parse_csv_floats(args.replay_scales) + if not replay_scales: + parser.error('--replay-scales must include at least one value') + if any(scale < 0 for scale in replay_scales): + parser.error('--replay-scales values must be non-negative') # Determine scenarios scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario] @@ -458,20 +711,28 @@ def main(): # Run benchmarks results = [] + use_scale_subdirs = len(replay_scales) > 1 or replay_scales[0] != 0 for scenario in scenarios: trace_path = Path(args.trace_dir) / trace_files[scenario] if trace_path.exists(): - result = run_benchmark( - str(trace_path), - str(Path(args.storage_dir) / scenario), - model_config, - args.max_requests, - args.max_pages, - args.page_size_tokens, - args.fsync_mode, - args.fsync_batch_size - ) - results.append(result) + for replay_scale in replay_scales: + run_dir = Path(args.storage_dir) / scenario + if use_scale_subdirs: + run_dir = run_dir / f"replay_{replay_scale:g}x" + result = run_benchmark( + str(trace_path), + str(run_dir), + model_config, + args.max_requests, + args.max_pages, + args.page_size_tokens, + args.fsync_mode, + args.fsync_batch_size, + args.threads, + replay_scale, + args.progress_interval + ) + results.append(result) else: print(f"Warning: Trace file not found: {trace_path}") diff --git a/benchmarks/storage_benchmark_v1/doc/README.md b/benchmarks/storage_benchmark_v1/doc/README.md index fa490aadff..8b53be660a 100644 --- a/benchmarks/storage_benchmark_v1/doc/README.md +++ b/benchmarks/storage_benchmark_v1/doc/README.md @@ -28,12 +28,47 @@ python benchmark.py --scenario conversation \ | `--max-pages` | `2000` | Maximum number of pages (creates modulo mapping if trace is larger) | | `--fsync-mode` | `none` | When to fsync: `none`, `batch`, `always`, or `end` | | `--fsync-batch-size` | `100` | Number of writes between fsync in batch mode | +| `--threads` | `1` | Number of benchmark client worker threads | +| `--replay-scales` | `0` | Comma-separated trace fast-forward speeds; `0` means unpaced | +| `--progress-interval` | `100` | Print progress every N requests; `0` disables per-request progress | + +### Replay Scale + +Use `--replay-scales` to run the same trace at different fast-forward speeds: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --replay-scales 1,2,4,8 +``` + +For example, `2` means 2x fast-forward and `8` means 8x fast-forward. `0` +preserves the old unpaced behavior. + +### Client Threads + +Use `--threads` to add benchmark client worker threads: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --threads 4 +``` + +With `--threads > 1`, each benchmark client thread uses an independent storage +file under `thread_N/data.bin`, similar to running multiple clients at the same +time. Final results aggregate the per-thread counters and latency samples. For +strict single-client trace-order read/write and hit-rate accounting, use +`--threads 1`. ## Output Format ### Progress Output -During execution, each request displays real-time statistics: +During execution, progress is printed every `--progress-interval` requests and +at the end of the run: ``` [ 10/12031] ids= 35 tokens= 18060 | QPS= 2.45 | R= 36 ( 22.01ms, 2435.2MB/s) | W= 963 ( 19.35ms, 2770.1MB/s) @@ -56,12 +91,26 @@ Fields: [General] Model: glm5 + Threads: 1 + Fast-forward: unpaced Requests: 12031 Tokens: 123456789 Total I/O Time: 245.123 s QPS: 49.07 Hit Rate: 3.25% +[Request Wall Latency] + Avg: 20.912 ms + P50: 19.654 ms + P95: 28.123 ms + P99: 34.987 ms + +[Request Storage I/O Latency] + Avg: 20.312 ms + P50: 18.987 ms + P95: 27.456 ms + P99: 33.210 ms + [Read Operations] Count: 390 Data Volume: 20919.62 MB @@ -90,6 +139,28 @@ Fields: Sync Count: 0 ``` +`Request Wall Latency` measures the benchmark client's wall-clock time spent +processing a request after replay pacing. `Request Storage I/O Latency` is the +sum of the request's page read/write latencies. Read/write operation latency is +reported per page operation. Percentile values use linear interpolation. + +## Measurement Notes + +- The default `--fsync-mode none` measures page-cache-backed write behavior. It + does not represent durable write latency. Use `--fsync-mode always`, `batch`, + or `end` when persistence cost is part of the benchmark target. +- `pread`/`pwrite` latency is measured from user space, so it can include page + cache effects, OS scheduling, and Python benchmark-client overhead. Treat the + reported latency as an observed storage-path latency, not raw device service + time. +- With `--threads > 1`, each thread replays the full trace as an independent + benchmark client with its own storage file. This is a multi-client drive test, + not parallel execution of one trace stream. +- For publication-quality numbers, use a fixed machine and storage device, + clear or isolate benchmark storage directories between runs, disable + per-request progress output with `--progress-interval 0`, and run multiple + trials before reporting stable statistics. + ## Modulo Mapping When the trace requires more pages than `--max-pages`, modulo mapping is enabled: diff --git a/benchmarks/storage_benchmark_v1/storage/disk.py b/benchmarks/storage_benchmark_v1/storage/disk.py index 286ff7510f..64089a6b44 100644 --- a/benchmarks/storage_benchmark_v1/storage/disk.py +++ b/benchmarks/storage_benchmark_v1/storage/disk.py @@ -18,11 +18,16 @@ def calc_percentiles(data): return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} import statistics sorted_data = sorted(data) - n = len(sorted_data) + def get_percentile(p): - idx = int(n * p / 100) - if idx >= n: idx = n - 1 - return sorted_data[idx] + if len(sorted_data) == 1: + return sorted_data[0] + rank = (len(sorted_data) - 1) * (p / 100) + lower = int(rank) + upper = min(lower + 1, len(sorted_data) - 1) + weight = rank - lower + return sorted_data[lower] * (1.0 - weight) + sorted_data[upper] * weight + return { 'avg_ms': statistics.mean(data), 'p50_ms': get_percentile(50), From 09c32f9b1c330d653942bb785e8c48cec5881b5a Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Wed, 8 Jul 2026 10:42:31 +0800 Subject: [PATCH 31/57] [TENT] Deadline-infeasible drop + degradation hook (RFC #2519 step 3) (#2764) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Deadline-infeasible drop + degradation hook (RFC #2519 step 3) Builds on step 2 (#2763 EDF ordering) and step 1 (#2618 deadline_ns + MLU observability). Adds the degradation layer: owners predicted to miss their deadline are dropped from dispatch instead of sent, and a local-decode signal is raised so the caller can recompute locally. - QueueLimits.mlu_local_threshold (θ_local, default 0 = disabled). - setDegradationPolicy(bandwidth_provider, hooks, now_provider): dependency- injected so the admission queue stays decoupled from the device-selection layer and remains unit-testable. now_provider defaults to steady_clock. - pickForDispatch gains an optional `dropped_owner_ids` out-param. When drop is enabled (θ_local > 0, deadline_aware, bandwidth provider set), an owner whose predicted MLU (= length/bw / (deadline - now)) reaches θ_local — or whose deadline is already past — is charged out of the outstanding accounting, marked terminal (CANCELED), reported in dropped_owner_ids, and triggers on_local_decode_suggested. Everything else dispatches as before. Interface only: no codec / local-decode body (those live in vLLM/SGLang; TENT only raises the signal), as scoped in the RFC. Strictly additive and fully opt-in — θ_local = 0 (default) means zero behavior change from step 2. 5 new unit tests (drop infeasible / keep feasible / expired-deadline drop / disabled-when-threshold-zero / no-drop-without-bandwidth-provider, incl. hook invocation and outstanding-accounting checks). Full admission_queue_test suite passes 21/21. Motivation (H20 / CX-7 RoCEv2, TENT backend): the MLU sweep in #2519 shows the 100us deadline tier sits at mean MLU 2.63 — well past a θ_local of ~1.5 — so under contention these transfers are provably infeasible and are exactly what step 3 would drop to local-decode rather than waste bandwidth on. Co-Authored-By: Claude Opus 4.8 * ci: re-trigger (tent-ci apt/sccache network flake, not a code failure) * ci: re-trigger build-flags (runner killed mid-build, resource flake #2611) --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 --- .../include/tent/runtime/admission_queue.h | 48 ++++++++- .../tent/src/runtime/admission_queue.cpp | 73 +++++++++++++- .../tent/tests/admission_queue_test.cpp | 98 +++++++++++++++++++ 3 files changed, 214 insertions(+), 5 deletions(-) diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index d5e970b15c..8746bcd464 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,14 @@ struct QueueLimits { // them. This only reorders selection within the existing capacity limits; // it does not admit/reject or otherwise change what gets dispatched. bool deadline_aware{false}; + // Opt-in deadline-infeasible drop (RFC #2519 step 3). Local-decode MLU + // threshold θ_local. 0 (default) disables drop entirely — behavior is the + // step-2 EDF ordering (or FIFO). When > 0 (e.g. 1.5) and a bandwidth + // provider is set, an owner whose predicted MLU + // (= predicted_transfer_time / remaining_window) reaches this threshold is + // dropped instead of dispatched, and on_local_decode_suggested is raised so + // the caller can recompute locally. Requires deadline_aware = true. + double mlu_local_threshold{0.0}; }; struct QueueOwnerInput { @@ -66,6 +75,23 @@ struct QueueSubmit { std::vector owners; }; +// RFC #2519 step 3: degradation signal raised when a transfer is predicted to +// miss its deadline and is dropped from dispatch. The bodies (compression / +// local recompute) live in the upper layer (vLLM/SGLang); TENT only raises the +// signal. No hook registered ⇒ the drop still happens but nothing is notified. +struct DegradationHooks { + std::function on_local_decode_suggested; +}; + +// Returns the predicted transfer bandwidth in bytes/second, or <= 0 if unknown +// (in which case the drop decision is skipped). Injected by the owner so the +// admission queue does not depend on the device-selection layer directly. +using BandwidthProvider = std::function; + +// Returns "now" as a steady-clock timestamp in nanoseconds, matching the units +// of Request.deadline_ns. Injectable so tests are deterministic. +using NowProvider = std::function; + // Runtime-private admission model. It is intentionally single-threaded; the // eventual TransferEngineImpl integration owns synchronization. class LocalTransferAdmissionQueue { @@ -82,8 +108,21 @@ class LocalTransferAdmissionQueue { Status tryAdmit(const QueueSubmit& submit, std::vector& admitted_owner_ids); - std::vector pickForDispatch(size_t max_owners, - size_t max_bytes); + // Returns the owners to dispatch. When step-3 drop is enabled + // (mlu_local_threshold > 0, deadline_aware, and a bandwidth provider set), + // owners predicted to miss their deadline are dropped: charged out of the + // outstanding accounting, marked terminal (CANCELED), appended to + // `dropped_owner_ids` (if non-null), and on_local_decode_suggested is + // raised. `dropped_owner_ids` is cleared on entry. + std::vector pickForDispatch( + size_t max_owners, size_t max_bytes, + std::vector* dropped_owner_ids = nullptr); + + // Install the step-3 degradation policy inputs. Optional; without it the + // queue never drops (default behavior). now defaults to steady_clock. + void setDegradationPolicy(BandwidthProvider bandwidth_provider, + DegradationHooks hooks, + NowProvider now_provider = nullptr); Status complete(QueueOwnerId owner_id, TransferStatusEnum terminal_status); @@ -124,6 +163,11 @@ class LocalTransferAdmissionQueue { size_t outstanding_bytes_{0}; size_t outstanding_user_owners_{0}; size_t outstanding_user_bytes_{0}; + + // RFC #2519 step 3 degradation policy (all optional / opt-in). + BandwidthProvider bandwidth_provider_; + DegradationHooks degradation_hooks_; + NowProvider now_provider_; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index e236e54a5c..e13fdb0bb7 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -15,8 +15,10 @@ #include "tent/runtime/admission_queue.h" #include +#include #include #include +#include namespace mooncake { namespace tent { @@ -213,8 +215,18 @@ Status LocalTransferAdmissionQueue::tryAdmit( return Status::OK(); } +void LocalTransferAdmissionQueue::setDegradationPolicy( + BandwidthProvider bandwidth_provider, DegradationHooks hooks, + NowProvider now_provider) { + bandwidth_provider_ = std::move(bandwidth_provider); + degradation_hooks_ = std::move(hooks); + now_provider_ = std::move(now_provider); +} + std::vector LocalTransferAdmissionQueue::pickForDispatch( - size_t max_owners, size_t max_bytes) { + size_t max_owners, size_t max_bytes, + std::vector* dropped_owner_ids) { + if (dropped_owner_ids) dropped_owner_ids->clear(); std::vector picked; if (max_owners == 0 || max_bytes == 0) return picked; @@ -222,8 +234,53 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( // deadline_aware, fifo_ is kept EDF-ordered at admission time (see // tryAdmit's ordered insert), so there is nothing to sort here — we just // consume from the front. This keeps the hot dispatch path O(picked) - // instead of re-sorting the whole queue (up to max_outstanding_owners) on - // every call. Default (deadline_aware == false) is plain FIFO. + // instead of re-sorting the whole queue on every call. Default + // (deadline_aware == false) is plain FIFO. + // + // RFC #2519 step 3 (opt-in): drop is active only when a positive threshold, + // deadline awareness, and a bandwidth provider are all present. + const bool drop_enabled = limits_.deadline_aware && + limits_.mlu_local_threshold > 0.0 && + static_cast(bandwidth_provider_); + const double bw_bps = drop_enabled ? bandwidth_provider_() : 0.0; + const uint64_t now_ns = + drop_enabled + ? (now_provider_ + ? now_provider_() + : static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now() + .time_since_epoch()) + .count())) + : 0; + + // Predicted MLU = predicted_transfer_time / remaining_window. Returns true + // if the owner is predicted to miss its deadline hard enough to drop. + auto shouldDrop = [&](const QueueOwner& owner) -> bool { + if (!drop_enabled || bw_bps <= 0.0) return false; + const uint64_t deadline_ns = owner.request.deadline_ns; + if (deadline_ns == 0) return false; // no deadline → never dropped + if (deadline_ns <= now_ns) return true; // already past → infeasible + const double window_s = (deadline_ns - now_ns) / 1e9; + const double predicted_time_s = owner.request.length / bw_bps; + const double mlu = predicted_time_s / window_s; + return mlu >= limits_.mlu_local_threshold; + }; + + auto dropOwner = [&](QueueOwnerId owner_id, QueueOwner& owner) { + owner.state = QueueState::Terminal; + owner.terminal_status = TransferStatusEnum::CANCELED; + --outstanding_owners_; + outstanding_bytes_ -= owner.request.length; + if (owner.kind == QueueOwnerKind::User) { + --outstanding_user_owners_; + outstanding_user_bytes_ -= owner.request.length; + } + if (dropped_owner_ids) dropped_owner_ids->push_back(owner_id); + if (degradation_hooks_.on_local_decode_suggested) { + degradation_hooks_.on_local_decode_suggested(owner.request); + } + }; size_t used_owners = 0; size_t used_bytes = 0; @@ -239,6 +296,16 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( continue; } + // Step 3: an owner predicted to miss its deadline is dropped (not + // dispatched) and does not consume the dispatch budget. Because the + // queue is EDF-ordered, later owners have looser deadlines, so we keep + // scanning rather than stopping. + if (shouldDrop(owner_it->second)) { + fifo_.pop_front(); + dropOwner(owner_id, owner_it->second); + continue; + } + const auto& owner = owner_it->second; const size_t remaining_bytes = max_bytes - used_bytes; if (owner.request.length > remaining_bytes) break; diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index f22edce8c3..083f36994c 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -497,6 +497,104 @@ TEST(AdmissionQueueTest, DeadlineAwareOrdersAcrossSeparateAdmits) { EXPECT_EQ(picked, expected); } +// --- RFC #2519 step 3: deadline-infeasible drop + degradation hook -------- + +// Helper: build a queue with deadline_aware + a θ_local, a fixed bandwidth, +// and a fixed "now" clock so MLU is deterministic. +QueueLimits step3Limits(double theta_local) { + QueueLimits limits{4, 1 << 20, 0, 0}; + limits.deadline_aware = true; + limits.mlu_local_threshold = theta_local; + return limits; +} + +TEST(AdmissionQueueTest, Step3DropsInfeasibleAndKeepsFeasible) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + // Fixed now = 1e9 ns; bandwidth = 1e9 B/s (so 16 B takes 16 ns). + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + // owner 1: window = 10 ns → 16 B / 1e9 = 16 ns → MLU 1.6 ≥ 1.5 → DROP. + // owner 2: window = 1e6 ns → MLU ~1.6e-5 → feasible → dispatch. + auto status = queue.tryAdmit( + makeSubmit(1, 2, + {makeOwnerWithDeadline(0, 16, 1'000'000'010), + makeOwnerWithDeadline(1, 16, 2'000'000'000)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + ASSERT_EQ(admitted_ids.size(), 2u); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + + const std::vector exp_pick{2}; + const std::vector exp_drop{1}; + EXPECT_EQ(picked, exp_pick); + EXPECT_EQ(dropped, exp_drop); + EXPECT_EQ(hook_calls, 1); + // Dropped owner is charged out of the outstanding accounting. + EXPECT_EQ(queue.outstandingOwners(), 1u); + EXPECT_EQ(queue.outstandingBytes(), 16u); +} + +TEST(AdmissionQueueTest, Step3DropsAlreadyExpiredDeadline) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{2'000'000'000}; }); + + std::vector admitted_ids; + // deadline 1e9 < now 2e9 → already past → dropped. + auto status = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1'000'000'000)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + EXPECT_TRUE(picked.empty()); + ASSERT_EQ(dropped.size(), 1u); + EXPECT_EQ(dropped[0], 1u); +} + +TEST(AdmissionQueueTest, Step3DisabledWhenThresholdZero) { + // θ_local = 0 (default off): even a hopeless deadline is dispatched, and + // the dropped vector stays empty — behavior is pure step-2 EDF. + LocalTransferAdmissionQueue queue(step3Limits(0.0)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + auto status = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1'000'000'001)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_EQ(picked[0], 1u); + EXPECT_TRUE(dropped.empty()); +} + +TEST(AdmissionQueueTest, Step3NoDropWithoutBandwidthProvider) { + // Threshold set but no bandwidth provider → cannot predict → never drops. + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + std::vector admitted_ids; + auto status = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1'000'000'001)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(dropped.empty()); +} + } // namespace } // namespace tent } // namespace mooncake From 454c349c5d600b300de319bda8e81056b2e8ffa5 Mon Sep 17 00:00:00 2001 From: "Guocheng(Eric) Song" Date: Wed, 8 Jul 2026 11:04:52 +0800 Subject: [PATCH 32/57] [Store] Fix S3 list objects pagination (#2778) --- mooncake-store/src/utils/s3_helper.cpp | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/mooncake-store/src/utils/s3_helper.cpp b/mooncake-store/src/utils/s3_helper.cpp index ce394a9904..fd2f9ead24 100644 --- a/mooncake-store/src/utils/s3_helper.cpp +++ b/mooncake-store/src/utils/s3_helper.cpp @@ -6,7 +6,7 @@ #include #include #include -#include +#include #include #include #include @@ -781,7 +781,7 @@ tl::expected S3Helper::ListObjectsWithPrefix( const std::string &prefix, std::vector &object_keys) { object_keys.clear(); - Aws::S3::Model::ListObjectsRequest request; + Aws::S3::Model::ListObjectsV2Request request; request.WithBucket(bucket_); request.WithPrefix(prefix); @@ -791,10 +791,10 @@ tl::expected S3Helper::ListObjectsWithPrefix( bool done = false; while (!done) { - auto outcome = s3_client_.ListObjects(request); + auto outcome = s3_client_.ListObjectsV2(request); if (!outcome.IsSuccess()) { return tl::make_unexpected(fmt::format( - "ListObjects error: {}", outcome.GetError().GetMessage())); + "ListObjectsV2 error: {}", outcome.GetError().GetMessage())); } const auto &result = outcome.GetResult(); @@ -806,8 +806,13 @@ tl::expected S3Helper::ListObjectsWithPrefix( // Check if there are more objects to fetch if (result.GetIsTruncated()) { - // Set marker to get next page - request.WithMarker(result.GetNextMarker()); + const auto &next_token = result.GetNextContinuationToken(); + if (next_token.empty()) { + return tl::make_unexpected( + "ListObjectsV2 error: truncated response missing next " + "continuation token"); + } + request.SetContinuationToken(next_token); } else { done = true; } From 83ff47f6fba457a42bf80d48d140f60cfc0d5ce7 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Wed, 8 Jul 2026 11:39:55 +0800 Subject: [PATCH 33/57] [Store] feat: add optional RFC #1527 KV events publisher on master (#2214) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(store): add optional RFC #1527 KV events publisher on master Implement an opt-in ZMQ publisher on mooncake_master that emits standardized KV cache events for Dynamo global KV indexer integration. - Add KvEventPublisher with async bounded queue and background worker - Publish stored/removed events on PutEnd, Remove, and eviction paths - Wire RFC #1527 envelope fields plus optional vLLM/SGLang compat aliases - Expose gflags/config toggles (enable_kv_events, bind endpoint, backend_id) - Add GET /kv_events/status on the master metrics HTTP server - Document Mooncake master publisher usage in indexer API design Co-authored-by: Teng Ma * docs(conductor): add KV event field provenance matrix (SGLang vs master vs register) Co-authored-by: Teng Ma * fix(store): address KV events PR review feedback - Fix msgpack map sizes with ComputeEventMapSize helper - Close ZMQ message parts on partial send failure - Emit per-medium removed events when memory replicas are evicted - Remove spurious removed event on PutRevoke before PutEnd - Use condition_variable::wait, drain full queue on shutdown, htobe64 - Drop bounded queue/drop-on-full to avoid indexer consistency gaps - Make libzmq optional via ENABLE_KV_EVENTS (stub when disabled) - Add missing condition_variable include and key_util for hash parsing Co-authored-by: Teng Ma * style(store): apply clang-format-20 to KV events changes Co-authored-by: Teng Ma * fix(store): fix KV events CI link failure when libzmq is absent Propagate ENABLE_KV_EVENTS=OFF to the parent CMake scope when libzmq is not found so kv_event_publisher_test uses the header stub consistently with mooncake_store. Add libzmq3-dev to dependencies.sh for full builds. Co-authored-by: Teng Ma * fix(store): link libzmq in Rust bindings when KV events enabled mooncake_store pulls in ZMQ symbols when libzmq is present; propagate the optional -lzmq link in build.rs using the same has_library pattern as uring/etcd. Co-authored-by: Teng Ma * fix(store): link libzmq in Go CGO flags for KV events Go integration tests link libmooncake_store.a statically and need -lzmq when the KV events publisher is compiled in. Co-authored-by: Teng Ma * ci: re-trigger build after runner disk-space failure Previous build (3.12) failed during Codecov upload with 'No space left on device' after all tests passed. No code changes. Co-authored-by: Teng Ma * fix(store): disable KV events compile flag when libzmq is absent When libzmq was not found, ENABLE_KV_EVENTS was only cleared in the parent CMake scope while the local value stayed ON. That defined MOONCAKE_ENABLE_KV_EVENTS=1 without linking kv_event_publisher.cpp, breaking mooncake_master on platforms like the Ascend CI image. Co-authored-by: Teng Ma * fix(store): address review comments — msgpack map_size, ZMQ leak, PutRevoke event, cmake guard - PutRevoke: emit PublishKvRemoved before erasing an invalidated object so that downstream indexers are notified when a revoked key is removed - cmake: upgrade ZMQ-not-found diagnostic from WARNING to FATAL_ERROR when ENABLE_KV_EVENTS is explicitly ON, since the user opted in - Verified: msgpack map_size values (18/15/15/13) are already correct after ComputeEventMapSize refactor in ac17fbc5 - Verified: ZMQ zmq_msg_close is already called on every send-failure path after the fix in ac17fbc5 Co-Authored-By: Claude Opus 4.6 * fix(cmake): default ENABLE_KV_EVENTS to OFF KV events require libzmq which is not available on all CI runners (e.g. Ascend). Default to OFF so builds succeed without it; users who want the feature opt in with -DENABLE_KV_EVENTS=ON. Co-Authored-By: Claude Opus 4.6 * fix(store): align KV events publisher with RFC #1527 spec - Set stored base_block_idx=0 so at least one placement field is present - Emit per-object tenant_id from master metadata, not only global config - Stop spurious removed events on PutRevoke (abort before PutEnd) - Publish removed on disk/NOF eviction paths with medium=disk - Pass explicit medium to eviction removed helper Co-authored-by: Teng Ma * fix(store): address PR #2214 review and CI issues - Fix clang-format on kv_event_publisher.h (CI Check code format) - Add optional object_key field (kv_events_emit_object_key, default on) for Dynamo matching on Mooncake store keys without decimal/0x seq_hash encoding - Publish events with empty seq_hashes when object_key is emitted but hash cannot be parsed; update indexer API docs and field matrix - Wire emit_object_key through master config, gflags, and master.json Co-authored-by: Teng Ma * style: clang-format master_admin_service.cpp Co-authored-by: Teng Ma * fix(store): emit per-block KV events without global semantic fields Align Mooncake master publisher with Dynamo KV events model: each event describes a pooled block (seq_hash/object_key, medium, tenant_id), not process-wide model/block_size/lora/dp_rank invariants. Omit unknown envelope fields (nil) and supply stream dimensions via indexer POST /register. Deprecate master flags that previously stamped every event. Co-authored-by: Teng Ma * feat(store): include group ID in KV events Signed-off-by: Ishan Dhanani --------- Signed-off-by: Ishan Dhanani Co-authored-by: Cursor Agent Co-authored-by: Teng Ma Co-authored-by: Claude Opus 4.6 Co-authored-by: Ishan Dhanani --- .github/workflows/ci.yml | 2 + dependencies.sh | 1 + .../design/conductor/indexer-api-design.md | 130 ++++++ mooncake-store/CMakeLists.txt | 3 + mooncake-store/conf/master.json | 10 +- mooncake-store/go/build.sh | 7 + mooncake-store/include/kv_event/key_util.h | 36 ++ .../include/kv_event/kv_event_config.h | 35 ++ .../include/kv_event/kv_event_publisher.h | 124 ++++++ mooncake-store/include/master_admin_service.h | 2 + mooncake-store/include/master_config.h | 98 +++++ mooncake-store/include/master_service.h | 24 ++ mooncake-store/include/rpc_service.h | 4 + mooncake-store/rust/build.rs | 1 + mooncake-store/src/CMakeLists.txt | 22 + .../src/kv_event/kv_event_publisher.cpp | 379 ++++++++++++++++++ mooncake-store/src/master.cpp | 126 ++++++ mooncake-store/src/master_admin_service.cpp | 30 ++ mooncake-store/src/master_service.cpp | 127 ++++++ mooncake-store/src/rpc_service.cpp | 8 + mooncake-store/tests/CMakeLists.txt | 15 + .../tests/kv_event_publisher_test.cpp | 218 ++++++++++ 22 files changed, 1401 insertions(+), 1 deletion(-) create mode 100644 mooncake-store/include/kv_event/key_util.h create mode 100644 mooncake-store/include/kv_event/kv_event_config.h create mode 100644 mooncake-store/include/kv_event/kv_event_publisher.h create mode 100644 mooncake-store/src/kv_event/kv_event_publisher.cpp create mode 100644 mooncake-store/tests/kv_event_publisher_test.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bb2e9925b7..345f31638f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -182,6 +182,8 @@ jobs: export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov -lxxhash -lyaml-cpp" # Link cudart if CUDA is available (needed for D2H staging in mooncake_store) if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi + # KV events publisher (optional; linked when libzmq is installed) + if ldconfig -p 2>/dev/null | grep -q libzmq; then export CGO_LDFLAGS="$CGO_LDFLAGS -lzmq"; fi ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/... kill $MASTER_PID 2>/dev/null || true shell: bash diff --git a/dependencies.sh b/dependencies.sh index bbbd82aa32..955d0494c0 100755 --- a/dependencies.sh +++ b/dependencies.sh @@ -175,6 +175,7 @@ if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then liburing-dev \ libjemalloc-dev \ libmsgpack-dev \ + libzmq3-dev \ libzstd-dev \ libasio-dev \ libxxhash-dev \ diff --git a/docs/source/design/conductor/indexer-api-design.md b/docs/source/design/conductor/indexer-api-design.md index e79367e6cd..c44af602fc 100644 --- a/docs/source/design/conductor/indexer-api-design.md +++ b/docs/source/design/conductor/indexer-api-design.md @@ -360,3 +360,133 @@ normalizes `BlockStored` and `BlockRemoved` into the internal prefix index. Registration metadata supplies fields such as `modelname`, `tenant_id`, `instance_id`, `block_size`, and `additionalsalt` when the engine event does not carry the full standardized envelope. + +### Mooncake Store master publisher + +`mooncake_master` can optionally publish RFC #1527 events when +`enable_kv_events=true`. The publisher binds a ZMQ PUB socket +(`kv_events_bind_endpoint`) and emits the same three-frame batch format used by +vLLM/SGLang: empty topic, big-endian sequence number, and a msgpack payload +`[timestamp, [events], dp_rank]`. + +**Per-block events, not global metadata.** Per the +[Dynamo KV Events for Custom Engines](https://docs.nvidia.com/dynamo/kv-managers/kv-events-for-custom-engines) +model, each event describes one or more **KV cache blocks** (`seq_hashes`, +`token_ids`, `parent_hash`, eviction hashes). The master emits **one event per +Mooncake object key** on `PutEnd` / `Remove` / eviction — each key is treated as +one pooled block. Block identity comes from the object key (`seq_hashes` when +the key is decimal/`0x` u64, else `object_key`) and per-object `tenant_id` / +`medium`. The master does **not** stamp process-wide `model_name`, `block_size`, +`lora_name`, or `dp_rank` on events; register those dimensions with the indexer +via `POST /register` (same as decoupled SGLang + storage pool deployments). + +Publisher-level config is limited to transport and stream identity: +`kv_events_bind_endpoint`, `kv_events_backend_id`, and optional compat flags +(`kv_events_emit_object_key`, `kv_events_emit_legacy_compat`). Legacy master +flags such as `kv_events_model_name` are retained for compatibility but are not +written into event payloads. + +Each event map uses RFC #1527 field names (`event_type`, `seq_hashes`, +`backend_id`, `medium`, and so on). When `kv_events_emit_object_key` is enabled +(default), the map also includes `object_key` with the Mooncake store key so +Dynamo and other consumers can match on `sha256` + Mooncake key format without +requiring decimal/`0x` `seq_hash` encoding. When `kv_events_emit_legacy_compat` +is enabled (default), the map also includes vLLM-compatible aliases such as +`type` and `block_hashes` so Dynamo relay mode can forward events without an +adapter. + +Object keys may encode the rolling `seq_hash` as a decimal or `0x`-prefixed +hex string; when `seq_hash` cannot be parsed, events are still published if +`kv_events_emit_object_key=true` (with an empty `seq_hashes` array). Configure +`backend_id` to identify the cache owner (for example a per-node storage +daemon) and register the bind endpoint with the indexer using publisher type +`Mooncake`. + +### Field provenance matrix (SGLang vs master vs indexer registration) + +In decoupled deployments (inference workers + Mooncake host/disk pool), the +global KV indexer merges **three sources of truth**. Use this table when +splitting publishers or writing PR/integration notes. + +**Legend** + +| Symbol | Meaning | +|---|---| +| **SGLang** | Inference engine ZMQ KV events (`BlockStored` / `BlockRemoved` / `AllBlocksCleared`) | +| **Master** | `mooncake_master` optional RFC #1527 publisher (`enable_kv_events`) | +| **Register** | Indexer HTTP `POST /register` (or CLI `--workers`) — not carried on the event wire | +| **S+M** | Either source may supply; must agree on value for the stream | +| **—** | Not applicable for that event type | + +#### Envelope and stream identity + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `event_id` | Yes | Yes | — | Each publisher maintains its own monotonic counter per stream. | +| `timestamp` | Yes | Yes | — | Informational only; not used for ordering. | +| `event_type` | Yes | Yes | — | `stored` / `removed` / `cleared`. | +| `model_name` | S+M | — | S+M | Register uses `modelname`. Engine events carry per-block context; master omits (nil). | +| `block_size` | Yes | — | Yes | Required for token↔block mapping. Register supplies for master publisher. | +| `additional_salt` | Yes | — | S+M | Register uses `additionalsalt`. Engine per-block; master omits (nil). | +| `lora_name` | Yes | — | S+M | Per-block on engine events; master has no adapter context. | +| `tenant_id` | S+M | Yes | Yes | Per-object on master events. Register default `default`. | +| `backend_id` | S+M | Yes | — | **Master**: storage daemon / pool owner. **SGLang**: often worker id; in decoupled mode prefer master=`daemon`, engine via **Register** `instance_id`. | +| `medium` | Yes | Partial | — | **SGLang**: `GPU`, `CPU_PINNED`, `DISK`, `EXTERNAL`, etc. **Master**: only `cpu` / `disk` (host/disk pool), never GPU. | +| `dp_rank` | Yes | — | Yes | Per-batch on engine ZMQ wire. Master batch trailer uses `0`; register dp_rank with indexer. | + +#### `stored` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `seq_hashes` | Yes | Conditional | — | **Required from SGLang** for correct prefix index. Master: single hash when key is decimal/`0x` u64; empty array when only `object_key` is used. | +| `object_key` | — | Yes | — | Mooncake store key (`kv_events_emit_object_key`, default on). Used by Dynamo for sha256+key matching. | +| `block_hashes` (legacy) | Yes | Conditional | — | Alias of `seq_hashes` when `kv_events_emit_legacy_compat` is enabled on master. | +| `parent_hash` | Yes | — | — | Radix parent link; master has no sequence tree. | +| `parent_block_hash` (legacy) | Yes | — | — | Same as `parent_hash`. | +| `base_block_idx` | Yes | Partial | — | Depth of first block in batch; master uses `0` for standalone pool blocks. | +| `token_ids` | Yes | — | — | Required for `/query` by tokens or hash recomputation when engine is non-standard. | +| `block_size` (in-event) | Yes | — | — | Per-block token count in SGLang `BlockStored`; master uses envelope-level config only. | + +#### `removed` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `seq_hashes` | Yes | Conditional | — | **Required** on wire for strict RFC consumers. Master emits one hash when parseable, else empty with `object_key`. | +| `base_block_idx` | Yes | — | — | Optional but recommended for observability. | + +#### `cleared` payload + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| (no extra fields) | — | — | — | Event is envelope-only. | +| `cleared` / `AllBlocksCleared` | Yes | — | — | Engine `reset()` / full cache flush. Master does not emit today. | + +#### Indexer / router plane (not in KV event JSON) + +| Field | SGLang | Master | Register | Notes | +|---|---|---|---|---| +| `instance_id` | — | — | Yes | Router-facing schedule target. Distinct from `backend_id`. | +| `endpoint` | — | — | Yes | ZMQ PUB to subscribe (SGLang or master bind address). | +| `replay_endpoint` | — | — | Yes | Optional gap replay (engine ROUTER). | +| `type` | — | — | Yes | Publisher kind: `vLLM`, `SGLang`, `Mooncake`, etc. | + +#### Recommended split for Dynamo global KV indexer + +```mermaid +flowchart LR + SGLang["SGLang ZMQ"] + Master["Mooncake master ZMQ"] + Reg["POST /register"] + Idx["Global KV indexer"] + + SGLang -->|"GPU + HiCache tiers
tokens, parent_hash, seq_hashes, lora"| Idx + Master -->|"Host/Disk pool
backend_id, medium=cpu|disk"| Idx + Reg -->|"instance_id, model, block_size"| Idx +``` + +| Capability | Primary source | +|---|---| +| GPU prefix hits, LoRA-aware hashes, parent chain, multi-block batches | **SGLang** | +| Pooled host/disk replica visibility | **Master** (if keys encode `seq_hash`) | +| Request routing target | **Register** (`instance_id`) | +| Tiered `/query` response (`gpu` / `cpu` / `disk`) | Merge **SGLang** + **Master** events (see RFC #1403) | diff --git a/mooncake-store/CMakeLists.txt b/mooncake-store/CMakeLists.txt index 5cf0f8a4ea..d03d82c67b 100644 --- a/mooncake-store/CMakeLists.txt +++ b/mooncake-store/CMakeLists.txt @@ -1,5 +1,8 @@ project(MooncakeStore VERSION 2.0.0) +option(ENABLE_KV_EVENTS + "Build master KV events ZMQ publisher (requires libzmq when ON)" OFF) + # Extract version components for C++ usage set(MOONCAKE_STORE_VERSION ${PROJECT_VERSION}) diff --git a/mooncake-store/conf/master.json b/mooncake-store/conf/master.json index a376e72fcb..5871146952 100644 --- a/mooncake-store/conf/master.json +++ b/mooncake-store/conf/master.json @@ -20,5 +20,13 @@ "client_live_ttl_sec": 60, "enable_http_metadata_server": false, "http_metadata_server_host": "0.0.0.0", - "http_metadata_server_port": 8080 + "http_metadata_server_port": 8080, + "enable_kv_events": false, + "kv_events_bind_endpoint": "tcp://0.0.0.0:5557", + "kv_events_model_name": "", + "kv_events_backend_id": "", + "kv_events_tenant_id": "default", + "kv_events_block_size": 0, + "kv_events_dp_rank": 0, + "kv_events_emit_object_key": true } diff --git a/mooncake-store/go/build.sh b/mooncake-store/go/build.sh index ac190e8b1a..7a751c9bc2 100755 --- a/mooncake-store/go/build.sh +++ b/mooncake-store/go/build.sh @@ -50,6 +50,13 @@ fi CGO_LDFLAGS+=" -luring" +# KV events publisher (optional; linked when libzmq is installed). +if ldconfig -p 2>/dev/null | grep -q libzmq \ + || [ -f /usr/lib/x86_64-linux-gnu/libzmq.so ] \ + || [ -f /usr/lib/libzmq.so ]; then + CGO_LDFLAGS+=" -lzmq" +fi + if [ "$USE_ETCD" = "ON" ]; then if [ "$USE_ETCD_LEGACY" = "ON" ]; then CGO_LDFLAGS+=" -letcd-cpp-api -lprotobuf -lgrpc++ -lgrpc" diff --git a/mooncake-store/include/kv_event/key_util.h b/mooncake-store/include/kv_event/key_util.h new file mode 100644 index 0000000000..6c90900ae8 --- /dev/null +++ b/mooncake-store/include/kv_event/key_util.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include +#include + +namespace mooncake { + +// Parses object keys encoded as decimal or 0x-prefixed hex u64 hashes. +inline std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + if (object_key.empty()) { + return std::nullopt; + } + try { + size_t idx = 0; + if (object_key.size() >= 2 && + (object_key[0] == '0' && + (object_key[1] == 'x' || object_key[1] == 'X'))) { + uint64_t value = std::stoull(object_key, &idx, 16); + if (idx == object_key.size()) { + return value; + } + return std::nullopt; + } + uint64_t value = std::stoull(object_key, &idx, 10); + if (idx == object_key.size()) { + return value; + } + } catch (const std::exception&) { + return std::nullopt; + } + return std::nullopt; +} + +} // namespace mooncake diff --git a/mooncake-store/include/kv_event/kv_event_config.h b/mooncake-store/include/kv_event/kv_event_config.h new file mode 100644 index 0000000000..9f255975d3 --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_config.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include + +namespace mooncake { + +// Publisher transport/identity for the optional KV Events ZMQ socket (RFC +// #1527). Semantic block fields (model_name, block_size, lora_name, +// parent_hash, token_ids, dp_rank) belong on each event payload, not here — see +// https://docs.nvidia.com/dynamo/kv-managers/kv-events-for-custom-engines +struct KvEventConfig { + bool enabled{false}; + // ZMQ PUB bind address, e.g. "tcp://0.0.0.0:5557". + std::string bind_endpoint; + // Identifies the cache owner stream (storage daemon / pool node). + std::string backend_id; + // Emit legacy vLLM/SGLang field names alongside RFC #1527 fields. + bool emit_legacy_compat_fields{true}; + // Emit Mooncake object_key for consumers that match on store key format. + bool emit_object_key{true}; + // Max pending events in the async publisher queue; oldest dropped when full. + uint32_t queue_capacity{65536}; + + // Deprecated: not stamped on events. Indexer registration supplies model, + // block_size, dp_rank, and hash namespace for the Mooncake publisher. + std::string model_name; + std::string tenant_id{"default"}; + std::string additional_salt; + std::string lora_name; + uint32_t block_size{0}; + uint32_t dp_rank{0}; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/kv_event/kv_event_publisher.h b/mooncake-store/include/kv_event/kv_event_publisher.h new file mode 100644 index 0000000000..39fc4af146 --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_publisher.h @@ -0,0 +1,124 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kv_event/kv_event_config.h" +#include "kv_event/key_util.h" + +namespace mooncake { + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +// Publishes standardized KV cache events (RFC #1527) over ZMQ for indexers. +class KvEventPublisher { + public: + explicit KvEventPublisher(KvEventConfig config); + ~KvEventPublisher(); + + KvEventPublisher(const KvEventPublisher&) = delete; + KvEventPublisher& operator=(const KvEventPublisher&) = delete; + + bool enabled() const { return config_.enabled; } + + // Non-blocking enqueue into a bounded queue; drops oldest when full. + // tenant_id empty defaults to "default" on the wire. + void PublishStored(const std::string& object_key, const std::string& medium, + const std::string& tenant_id = "", + const std::string& group_id = ""); + void PublishRemoved(const std::string& object_key, + const std::string& medium, + const std::string& tenant_id = "", + const std::string& group_id = ""); + + struct Stats { + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; + }; + Stats GetStats() const; + + static std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + return mooncake::ParseSeqHashFromObjectKey(object_key); + } + + private: + enum class EventKind { kStored, kRemoved }; + + struct PendingEvent { + EventKind kind; + std::string object_key; + std::string medium; + std::string tenant_id; + std::string group_id; + }; + + void Enqueue(PendingEvent event); + void WorkerLoop(); + void PublishBatch(const std::vector& batch); + void DrainRemainingQueue(std::vector& batch); + + KvEventConfig config_; + void* zmq_context_{nullptr}; + void* zmq_socket_{nullptr}; + + mutable std::mutex queue_mutex_; + std::deque queue_; + std::condition_variable queue_cv_; + std::thread worker_; + std::atomic stop_{false}; + + std::atomic next_event_id_{1}; + std::atomic next_zmq_sequence_{1}; + + std::atomic published_batches_{0}; + std::atomic published_events_{0}; + std::atomic dropped_events_{0}; + std::atomic skipped_unparsed_keys_{0}; +}; + +#else + +// Stub when mooncake_store is built without libzmq (ENABLE_KV_EVENTS=OFF). +class KvEventPublisher { + public: + explicit KvEventPublisher(KvEventConfig config) + : config_(std::move(config)) {} + + bool enabled() const { return false; } + + void PublishStored(const std::string&, const std::string&, + const std::string& = "", const std::string& = "") {} + void PublishRemoved(const std::string&, const std::string&, + const std::string& = "", const std::string& = "") {} + + struct Stats { + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; + }; + Stats GetStats() const { return {}; } + + static std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + return mooncake::ParseSeqHashFromObjectKey(object_key); + } + + private: + KvEventConfig config_; +}; + +#endif + +} // namespace mooncake diff --git a/mooncake-store/include/master_admin_service.h b/mooncake-store/include/master_admin_service.h index 1898cb43ef..5ced911819 100644 --- a/mooncake-store/include/master_admin_service.h +++ b/mooncake-store/include/master_admin_service.h @@ -93,6 +93,8 @@ class MasterAdminServer { coro_http::coro_http_response& resp); void HandleBatchQueryKeys(coro_http::coro_http_request& req, coro_http::coro_http_response& resp); + void HandleKvEventsStatus(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); void HandleGetTenantQuotas(coro_http::coro_http_request& req, coro_http::coro_http_response& resp); void HandleUpsertTenantQuota(coro_http::coro_http_request& req, diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 0275c57e0f..fd893953f8 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -132,6 +132,20 @@ struct MasterConfig { // liveness window. Default 1 is conservative; small-object or RDMA- // rich clusters may safely raise it. uint32_t promotion_max_per_heartbeat = 1; + + // KV Events publisher (RFC #1527) for cache-aware indexers. + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; }; class MasterServiceSupervisorConfig { @@ -211,6 +225,18 @@ class MasterServiceSupervisorConfig { uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; // Pod identity for K8s label-based routing std::string pod_name; @@ -253,6 +279,18 @@ class MasterServiceSupervisorConfig { promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; rpc_port = static_cast(config.rpc_port); rpc_thread_num = static_cast(config.rpc_thread_num); @@ -403,6 +441,18 @@ class WrappedMasterServiceConfig { uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string cluster_id = DEFAULT_CLUSTER_ID; @@ -476,6 +526,18 @@ class WrappedMasterServiceConfig { promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, @@ -573,6 +635,18 @@ class WrappedMasterServiceConfig { promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, @@ -983,6 +1057,18 @@ class MasterServiceConfig { uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string cluster_id = DEFAULT_CLUSTER_ID; @@ -1052,6 +1138,18 @@ class MasterServiceConfig { promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = config.ha_backend_connstring; cluster_id = config.cluster_id; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index d7823d2468..661109fd51 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -36,6 +36,7 @@ #include "ha/ha_types.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "task_manager.h" +#include "kv_event/kv_event_publisher.h" namespace mooncake { namespace ha { @@ -288,6 +289,9 @@ class MasterService { std::unordered_map, boost::hash>, ErrorCode>; + bool KvEventsEnabled() const; + KvEventPublisher::Stats GetKvEventStats() const; + /** * @brief Batch clear KV cache replicas for specified object keys. * @param object_keys Vector of object key strings to clear. @@ -2130,6 +2134,26 @@ class MasterService { std::mutex job_mutex_; std::unordered_map, boost::hash> drain_jobs_ GUARDED_BY(job_mutex_); + + std::unique_ptr kv_event_publisher_; + + static KvEventConfig BuildKvEventConfig(const MasterServiceConfig& config); + static std::string MediumForReplicaType(ReplicaType replica_type); + static std::string MediumForMetadata(const ObjectMetadata& metadata); + void PublishKvStored(const std::string& key, ReplicaType replica_type, + const ObjectMetadata& metadata, + const std::string& tenant_id); + void PublishKvRemoved(const std::string& key, + const ObjectMetadata& metadata, + const std::string& tenant_id); + void PublishKvRemoved(const std::string& key, const std::string& medium, + const std::string& tenant_id, + const std::string& group_id); + void PublishKvRemovedAfterEvict(const std::string& key, + uint64_t freed_bytes, + const std::string& medium, + const ObjectMetadata& metadata, + const std::string& tenant_id); }; } // namespace mooncake diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index f54f2ec0dd..8e2cfa4895 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -12,6 +12,7 @@ #include "types.h" #include "rpc_types.h" #include "master_config.h" +#include "kv_event/kv_event_publisher.h" #include "segment.h" namespace mooncake { @@ -299,6 +300,9 @@ class WrappedMasterService { const UUID& client_id, const std::vector& keys, const std::string& tenant_id, ReplicaType replica_type); + bool KvEventsEnabled() const; + KvEventPublisher::Stats GetKvEventStats() const; + private: MasterService master_service_; }; diff --git a/mooncake-store/rust/build.rs b/mooncake-store/rust/build.rs index edf3710ba0..f8225e6818 100644 --- a/mooncake-store/rust/build.rs +++ b/mooncake-store/rust/build.rs @@ -363,6 +363,7 @@ fn main() { ("cudart", &["cudart"]), ("mlx5", &["mlx5"]), // IBGDA device transport (mlx5 DevX) pulled into transfer_engine, CUDA-only ("uring", &["uring"]), + ("zmq", &["zmq"]), ] { if has_library(&search_dirs, candidates) { println!("cargo:rustc-link-lib={link_name}"); diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 9caf538b7b..f85b42c280 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -114,6 +114,25 @@ find_library( XXHASH_LIBRARY NAMES xxhash libxxhash PATHS /usr/lib /usr/local/lib /usr/lib64) +set(KV_EVENTS_ZMQ_FOUND FALSE) +if(ENABLE_KV_EVENTS) + find_library(ZMQ_LIBRARY NAMES zmq libzmq PATHS /usr/lib /usr/local/lib + /usr/lib64) + find_path(ZMQ_INCLUDE_DIR NAMES zmq.h PATHS /usr/include /usr/local/include) + if(ZMQ_INCLUDE_DIR AND ZMQ_LIBRARY) + message(STATUS "Found ZMQ: include=${ZMQ_INCLUDE_DIR} lib=${ZMQ_LIBRARY}") + list(APPEND MASTER_EXTRA_INCS ${ZMQ_INCLUDE_DIR}) + list(APPEND EXTRA_LIBS ${ZMQ_LIBRARY}) + list(APPEND MOONCAKE_STORE_SOURCES kv_event/kv_event_publisher.cpp) + set(KV_EVENTS_ZMQ_FOUND TRUE) + else() + message( + FATAL_ERROR + "ENABLE_KV_EVENTS is ON but libzmq was not found. " + "Install libzmq3-dev or pass -DENABLE_KV_EVENTS=OFF to configure without ZMQ.") + endif() +endif() + if(XXHASH_INCLUDE_DIR AND XXHASH_LIBRARY) message( STATUS "Found xxHash: include=${XXHASH_INCLUDE_DIR} lib=${XXHASH_LIBRARY}") @@ -246,6 +265,9 @@ endif() # The cache_allocator library include_directories(${Python3_INCLUDE_DIRS}) add_library(mooncake_store ${MOONCAKE_STORE_SOURCES}) +if(KV_EVENTS_ZMQ_FOUND) + target_compile_definitions(mooncake_store PRIVATE MOONCAKE_ENABLE_KV_EVENTS=1) +endif() target_include_directories(mooncake_store PUBLIC ${XXHASH_INCLUDE_DIR}) if(USE_NOF) target_include_directories(mooncake_store PRIVATE ${SPDK_INCLUDE_DIR} diff --git a/mooncake-store/src/kv_event/kv_event_publisher.cpp b/mooncake-store/src/kv_event/kv_event_publisher.cpp new file mode 100644 index 0000000000..a819a9bf0f --- /dev/null +++ b/mooncake-store/src/kv_event/kv_event_publisher.cpp @@ -0,0 +1,379 @@ +#include "kv_event/kv_event_publisher.h" + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +#include +#include +#include + +#include +#include +#include +#include + +namespace mooncake { +namespace { + +constexpr int kZmqSendHwm = 10000; +constexpr size_t kMaxBatchSize = 64; + +int64_t CurrentUnixTimeMs() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} + +void PackOptionalString(msgpack::packer& packer, + const std::string& value) { + if (value.empty()) { + packer.pack_nil(); + } else { + packer.pack(value); + } +} + +void PackOptionalU32(msgpack::packer& packer, uint32_t value, + bool has_value) { + if (!has_value) { + packer.pack_nil(); + } else { + packer.pack(value); + } +} + +size_t ComputeEventMapSize(bool is_stored, bool emit_legacy, + bool emit_object_key) { + // Base envelope: event_id, timestamp, event_type, model_name, block_size, + // additional_salt, lora_name, tenant_id, backend_id, medium, dp_rank, + // seq_hashes, group_id. + constexpr size_t kBaseFields = 13; + size_t map_size = kBaseFields; + if (emit_legacy) { + map_size += 2; // type, block_hashes + } + if (emit_object_key) { + map_size += 1; // object_key + } + if (is_stored) { + map_size += 3; // base_block_idx, parent_hash, token_ids + if (emit_legacy) { + map_size += 1; // parent_block_hash + } + } else { + map_size += 1; // base_block_idx + } + return map_size; +} + +} // namespace + +KvEventPublisher::KvEventPublisher(KvEventConfig config) + : config_(std::move(config)) { + if (!config_.enabled) { + return; + } + if (config_.bind_endpoint.empty()) { + LOG(ERROR) << "kv_events enabled but bind_endpoint is empty"; + config_.enabled = false; + return; + } + if (config_.backend_id.empty()) { + LOG(ERROR) << "kv_events enabled but backend_id is empty"; + config_.enabled = false; + return; + } + + zmq_context_ = zmq_ctx_new(); + if (!zmq_context_) { + LOG(ERROR) << "kv_events: failed to create ZMQ context"; + config_.enabled = false; + return; + } + zmq_socket_ = zmq_socket(zmq_context_, ZMQ_PUB); + if (!zmq_socket_) { + LOG(ERROR) << "kv_events: failed to create ZMQ PUB socket: " + << zmq_strerror(zmq_errno()); + zmq_ctx_destroy(zmq_context_); + zmq_context_ = nullptr; + config_.enabled = false; + return; + } + int hwm = kZmqSendHwm; + zmq_setsockopt(zmq_socket_, ZMQ_SNDHWM, &hwm, sizeof(hwm)); + int linger_ms = 0; + zmq_setsockopt(zmq_socket_, ZMQ_LINGER, &linger_ms, sizeof(linger_ms)); + + if (zmq_bind(zmq_socket_, config_.bind_endpoint.c_str()) != 0) { + LOG(ERROR) << "kv_events: zmq_bind failed for " << config_.bind_endpoint + << ": " << zmq_strerror(zmq_errno()); + zmq_close(zmq_socket_); + zmq_ctx_destroy(zmq_context_); + zmq_socket_ = nullptr; + zmq_context_ = nullptr; + config_.enabled = false; + return; + } + + worker_ = std::thread(&KvEventPublisher::WorkerLoop, this); + LOG(INFO) << "kv_events publisher enabled on " << config_.bind_endpoint + << " backend_id=" << config_.backend_id; +} + +KvEventPublisher::~KvEventPublisher() { + if (!config_.enabled) { + return; + } + stop_.store(true); + queue_cv_.notify_all(); + if (worker_.joinable()) { + worker_.join(); + } + if (zmq_socket_) { + zmq_close(zmq_socket_); + zmq_socket_ = nullptr; + } + if (zmq_context_) { + zmq_ctx_destroy(zmq_context_); + zmq_context_ = nullptr; + } +} + +void KvEventPublisher::PublishStored(const std::string& object_key, + const std::string& medium, + const std::string& tenant_id, + const std::string& group_id) { + if (!config_.enabled) { + return; + } + Enqueue(PendingEvent{EventKind::kStored, object_key, medium, tenant_id, + group_id}); +} + +void KvEventPublisher::PublishRemoved(const std::string& object_key, + const std::string& medium, + const std::string& tenant_id, + const std::string& group_id) { + if (!config_.enabled) { + return; + } + Enqueue(PendingEvent{EventKind::kRemoved, object_key, medium, tenant_id, + group_id}); +} + +KvEventPublisher::Stats KvEventPublisher::GetStats() const { + Stats stats; + stats.published_batches = published_batches_.load(); + stats.published_events = published_events_.load(); + stats.dropped_events = dropped_events_.load(); + stats.skipped_unparsed_keys = skipped_unparsed_keys_.load(); + return stats; +} + +void KvEventPublisher::Enqueue(PendingEvent event) { + { + std::lock_guard lock(queue_mutex_); + if (config_.queue_capacity > 0 && + queue_.size() >= config_.queue_capacity) { + queue_.pop_front(); + dropped_events_.fetch_add(1, std::memory_order_relaxed); + // Reserve a ZMQ sequence gap so consumers can detect loss. + next_zmq_sequence_.fetch_add(1, std::memory_order_relaxed); + } + queue_.push_back(std::move(event)); + } + queue_cv_.notify_one(); +} + +void KvEventPublisher::DrainRemainingQueue(std::vector& batch) { + while (true) { + batch.clear(); + { + std::lock_guard lock(queue_mutex_); + if (queue_.empty()) { + break; + } + while (!queue_.empty() && batch.size() < kMaxBatchSize) { + batch.push_back(std::move(queue_.front())); + queue_.pop_front(); + } + } + PublishBatch(batch); + } +} + +void KvEventPublisher::WorkerLoop() { + std::vector batch; + batch.reserve(kMaxBatchSize); + while (!stop_.load()) { + { + std::unique_lock lock(queue_mutex_); + queue_cv_.wait(lock, + [this] { return stop_.load() || !queue_.empty(); }); + while (!queue_.empty() && batch.size() < kMaxBatchSize) { + batch.push_back(std::move(queue_.front())); + queue_.pop_front(); + } + } + if (!batch.empty()) { + PublishBatch(batch); + batch.clear(); + } + } + DrainRemainingQueue(batch); +} + +void KvEventPublisher::PublishBatch(const std::vector& batch) { + struct EncodedEvent { + PendingEvent pending; + std::optional seq_hash; + uint64_t event_id{0}; + }; + std::vector encoded; + encoded.reserve(batch.size()); + for (const auto& pending : batch) { + const auto seq_hash = ParseSeqHashFromObjectKey(pending.object_key); + if (!seq_hash.has_value()) { + if (!config_.emit_object_key || pending.object_key.empty()) { + skipped_unparsed_keys_.fetch_add(1, std::memory_order_relaxed); + continue; + } + skipped_unparsed_keys_.fetch_add(1, std::memory_order_relaxed); + } + encoded.push_back(EncodedEvent{ + pending, seq_hash, + next_event_id_.fetch_add(1, std::memory_order_relaxed)}); + } + if (encoded.empty()) { + return; + } + + msgpack::sbuffer payload_buffer; + msgpack::packer packer(&payload_buffer); + + const int64_t timestamp_ms = CurrentUnixTimeMs(); + + packer.pack_array(3); + packer.pack(timestamp_ms); + + packer.pack_array(encoded.size()); + for (const auto& item : encoded) { + const bool is_stored = item.pending.kind == EventKind::kStored; + const char* rfc_type = is_stored ? "stored" : "removed"; + const char* legacy_type = is_stored ? "BlockStored" : "BlockRemoved"; + const std::string& tenant_id = + item.pending.tenant_id.empty() ? "default" : item.pending.tenant_id; + + const size_t map_size = + ComputeEventMapSize(is_stored, config_.emit_legacy_compat_fields, + config_.emit_object_key); + + packer.pack_map(map_size); + packer.pack("event_id"); + packer.pack(item.event_id); + packer.pack("timestamp"); + packer.pack(timestamp_ms); + packer.pack("event_type"); + packer.pack(rfc_type); + if (config_.emit_legacy_compat_fields) { + packer.pack("type"); + packer.pack(legacy_type); + } + // Per-block envelope fields unknown to the storage pool are omitted + // (nil). Indexer registration supplies model/block_size/dp_rank. + packer.pack("model_name"); + packer.pack_nil(); + packer.pack("block_size"); + packer.pack_nil(); + packer.pack("additional_salt"); + packer.pack_nil(); + packer.pack("lora_name"); + packer.pack_nil(); + packer.pack("tenant_id"); + packer.pack(tenant_id); + packer.pack("backend_id"); + packer.pack(config_.backend_id); + packer.pack("group_id"); + PackOptionalString(packer, item.pending.group_id); + packer.pack("medium"); + PackOptionalString(packer, item.pending.medium); + packer.pack("dp_rank"); + packer.pack_nil(); + + if (config_.emit_object_key) { + packer.pack("object_key"); + packer.pack(item.pending.object_key); + } + + packer.pack("seq_hashes"); + if (item.seq_hash.has_value()) { + packer.pack_array(1); + packer.pack(item.seq_hash.value()); + } else { + packer.pack_array(0); + } + + if (config_.emit_legacy_compat_fields && item.seq_hash.has_value()) { + packer.pack("block_hashes"); + packer.pack_array(1); + packer.pack(static_cast(item.seq_hash.value())); + } else if (config_.emit_legacy_compat_fields) { + packer.pack("block_hashes"); + packer.pack_array(0); + } + + if (is_stored) { + // Master keys are standalone pool blocks; depth 0 satisfies RFC + // #1527 requirement that base_block_idx or parent_hash be present. + packer.pack("base_block_idx"); + packer.pack(static_cast(0)); + packer.pack("parent_hash"); + packer.pack_nil(); + packer.pack("token_ids"); + packer.pack_nil(); + if (config_.emit_legacy_compat_fields) { + packer.pack("parent_block_hash"); + packer.pack_nil(); + } + } else { + packer.pack("base_block_idx"); + packer.pack_nil(); + } + } + + // Batch-level dp_rank; storage pool has no DP context (0). + packer.pack(static_cast(0)); + + const uint64_t seq = next_zmq_sequence_.fetch_add(1); + const uint64_t seq_be = htobe64(seq); + + zmq_msg_t topic_msg; + zmq_msg_t seq_msg; + zmq_msg_t payload_msg; + zmq_msg_init_size(&topic_msg, 0); + zmq_msg_init_size(&seq_msg, sizeof(seq_be)); + std::memcpy(zmq_msg_data(&seq_msg), &seq_be, sizeof(seq_be)); + zmq_msg_init_size(&payload_msg, payload_buffer.size()); + std::memcpy(zmq_msg_data(&payload_msg), payload_buffer.data(), + payload_buffer.size()); + + const int rc_topic = zmq_sendmsg(zmq_socket_, &topic_msg, ZMQ_SNDMORE); + const int rc_seq = + (rc_topic >= 0) ? zmq_sendmsg(zmq_socket_, &seq_msg, ZMQ_SNDMORE) : -1; + const int rc_payload = + (rc_seq >= 0) ? zmq_sendmsg(zmq_socket_, &payload_msg, 0) : -1; + + zmq_msg_close(&topic_msg); + zmq_msg_close(&seq_msg); + zmq_msg_close(&payload_msg); + + if (rc_topic >= 0 && rc_seq >= 0 && rc_payload >= 0) { + published_batches_.fetch_add(1, std::memory_order_relaxed); + published_events_.fetch_add(encoded.size(), std::memory_order_relaxed); + } else { + dropped_events_.fetch_add(encoded.size(), std::memory_order_relaxed); + } +} + +} // namespace mooncake + +#endif // MOONCAKE_ENABLE_KV_EVENTS diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 9174bd97ee..6092ca83a3 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -225,6 +225,30 @@ DEFINE_uint32(promotion_max_per_heartbeat, 1, "SSD-read + RDMA-write on the client; serializing them avoids " "blocking past the client-liveness window. Default 1 is " "conservative."); +DEFINE_bool(enable_kv_events, false, + "Enable RFC #1527 KV cache event publisher over ZMQ"); +DEFINE_string(kv_events_bind_endpoint, "", + "ZMQ PUB bind endpoint for KV events, e.g. tcp://0.0.0.0:5557"); +DEFINE_string(kv_events_model_name, "", + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_string(kv_events_backend_id, "", + "backend_id for published KV events (cache owner identity)"); +DEFINE_string(kv_events_tenant_id, "default", + "Deprecated: tenant_id comes from each object on events"); +DEFINE_string(kv_events_additional_salt, "", + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_string(kv_events_lora_name, "", + "Deprecated: not emitted on events (no LoRA context in master)"); +DEFINE_uint32(kv_events_block_size, 0, + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_uint32(kv_events_dp_rank, 0, + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_bool(kv_events_emit_legacy_compat, true, + "Include vLLM/SGLang-compatible type/block_hashes fields"); +DEFINE_bool(kv_events_emit_object_key, true, + "Include Mooncake object_key in published KV events"); +DEFINE_uint32(kv_events_queue_capacity, 65536, + "Deprecated; ignored (event queue is unbounded)"); DEFINE_string(ha_backend_type, "etcd", "HA backend type, e.g. etcd | redis | k8s"); DEFINE_string(ha_backend_connstring, "", @@ -482,6 +506,41 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetUInt32("promotion_max_per_heartbeat", &master_config.promotion_max_per_heartbeat, FLAGS_promotion_max_per_heartbeat); + default_config.GetBool("enable_kv_events", &master_config.enable_kv_events, + FLAGS_enable_kv_events); + default_config.GetString("kv_events_bind_endpoint", + &master_config.kv_events_bind_endpoint, + FLAGS_kv_events_bind_endpoint); + default_config.GetString("kv_events_model_name", + &master_config.kv_events_model_name, + FLAGS_kv_events_model_name); + default_config.GetString("kv_events_backend_id", + &master_config.kv_events_backend_id, + FLAGS_kv_events_backend_id); + default_config.GetString("kv_events_tenant_id", + &master_config.kv_events_tenant_id, + FLAGS_kv_events_tenant_id); + default_config.GetString("kv_events_additional_salt", + &master_config.kv_events_additional_salt, + FLAGS_kv_events_additional_salt); + default_config.GetString("kv_events_lora_name", + &master_config.kv_events_lora_name, + FLAGS_kv_events_lora_name); + default_config.GetUInt32("kv_events_block_size", + &master_config.kv_events_block_size, + FLAGS_kv_events_block_size); + default_config.GetUInt32("kv_events_dp_rank", + &master_config.kv_events_dp_rank, + FLAGS_kv_events_dp_rank); + default_config.GetBool("kv_events_emit_legacy_compat", + &master_config.kv_events_emit_legacy_compat, + FLAGS_kv_events_emit_legacy_compat); + default_config.GetBool("kv_events_emit_object_key", + &master_config.kv_events_emit_object_key, + FLAGS_kv_events_emit_object_key); + default_config.GetUInt32("kv_events_queue_capacity", + &master_config.kv_events_queue_capacity, + FLAGS_kv_events_queue_capacity); default_config.GetString("ha_backend_type", &master_config.ha_backend_type, FLAGS_ha_backend_type); default_config.GetString("ha_backend_connstring", @@ -795,6 +854,70 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, master_config.promotion_max_per_heartbeat = FLAGS_promotion_max_per_heartbeat; } + if ((google::GetCommandLineFlagInfo("enable_kv_events", &info) && + !info.is_default) || + !conf_set) { + master_config.enable_kv_events = FLAGS_enable_kv_events; + } + if ((google::GetCommandLineFlagInfo("kv_events_bind_endpoint", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_bind_endpoint = FLAGS_kv_events_bind_endpoint; + } + if ((google::GetCommandLineFlagInfo("kv_events_model_name", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_model_name = FLAGS_kv_events_model_name; + } + if ((google::GetCommandLineFlagInfo("kv_events_backend_id", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_backend_id = FLAGS_kv_events_backend_id; + } + if ((google::GetCommandLineFlagInfo("kv_events_tenant_id", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_tenant_id = FLAGS_kv_events_tenant_id; + } + if ((google::GetCommandLineFlagInfo("kv_events_additional_salt", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_additional_salt = + FLAGS_kv_events_additional_salt; + } + if ((google::GetCommandLineFlagInfo("kv_events_lora_name", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_lora_name = FLAGS_kv_events_lora_name; + } + if ((google::GetCommandLineFlagInfo("kv_events_block_size", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_block_size = FLAGS_kv_events_block_size; + } + if ((google::GetCommandLineFlagInfo("kv_events_dp_rank", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_dp_rank = FLAGS_kv_events_dp_rank; + } + if ((google::GetCommandLineFlagInfo("kv_events_emit_legacy_compat", + &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_emit_legacy_compat = + FLAGS_kv_events_emit_legacy_compat; + } + if ((google::GetCommandLineFlagInfo("kv_events_emit_object_key", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_emit_object_key = + FLAGS_kv_events_emit_object_key; + } + if ((google::GetCommandLineFlagInfo("kv_events_queue_capacity", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_queue_capacity = FLAGS_kv_events_queue_capacity; + } // Clamp promotion_admission_threshold into the sketch counter's // representable range. The CountMinSketch uses 8-bit saturating // counters (max 255) so any threshold beyond that would silently @@ -1254,6 +1377,9 @@ int main(int argc, char* argv[]) { << master_config.eviction_high_watermark_ratio << ", enable_ha=" << master_config.enable_ha << ", enable_offload=" << master_config.enable_offload + << ", enable_kv_events=" << master_config.enable_kv_events + << ", kv_events_bind_endpoint=" << master_config.kv_events_bind_endpoint + << ", kv_events_backend_id=" << master_config.kv_events_backend_id << ", offload_on_evict=" << master_config.offload_on_evict << ", offload_force_evict=" << master_config.offload_force_evict << ", offloading_queue_limit=" << master_config.offloading_queue_limit diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp index cb221af5de..798ca55346 100644 --- a/mooncake-store/src/master_admin_service.cpp +++ b/mooncake-store/src/master_admin_service.cpp @@ -562,6 +562,31 @@ void MasterAdminServer::HandleHaStatus(coro_http::coro_http_request&, ha::MasterRuntimeStateToString(snapshot.state)); } +struct HttpKvEventsStatusResponse { + bool enabled{false}; + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; +}; +YLT_REFL(HttpKvEventsStatusResponse, enabled, published_batches, + published_events, dropped_events, skipped_unparsed_keys); + +void MasterAdminServer::HandleKvEventsStatus( + coro_http::coro_http_request&, coro_http::coro_http_response& resp) { + WithActiveService( + resp, [&](const std::shared_ptr& service) { + const auto stats = service->GetKvEventStats(); + HttpKvEventsStatusResponse payload; + payload.enabled = service->KvEventsEnabled(); + payload.published_batches = stats.published_batches; + payload.published_events = stats.published_events; + payload.dropped_events = stats.dropped_events; + payload.skipped_unparsed_keys = stats.skipped_unparsed_keys; + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); +} + void MasterAdminServer::HandleQueryKey(coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { WithActiveService(resp, [&](auto service) { @@ -1156,6 +1181,11 @@ void MasterAdminServer::RegisterHandler() { "/ha_status", [this](coro_http_request& req, coro_http_response& resp) { HandleHaStatus(req, resp); }); + http_server_.set_http_handler( + "/kv_events/status", + [this](coro_http_request& req, coro_http_response& resp) { + HandleKvEventsStatus(req, resp); + }); http_server_.set_http_handler( "/leader", [this](coro_http_request& req, coro_http_response& resp) { HandleLeader(req, resp); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 5f55c2ed76..f3c2403dc5 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -44,6 +44,7 @@ #include "utils/zstd_util.h" #include "utils/file_util.h" #include "utils.h" +#include "kv_event/kv_event_config.h" namespace mooncake { @@ -409,6 +410,9 @@ MasterService::MasterService(const MasterServiceConfig& config) << ")"; } + kv_event_publisher_ = + std::make_unique(BuildKvEventConfig(config)); + eviction_running_ = true; eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this); VLOG(1) << "action=start_eviction_thread"; @@ -3288,6 +3292,7 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key, // at beginning. 2. If this object has soft pin enabled, set it to be soft // pinned. metadata.GrantLease(0, default_kv_soft_pin_ttl_); + PublishKvStored(key, replica_type, metadata, tenant_id); return {}; } @@ -3877,6 +3882,7 @@ auto MasterService::EvictDiskReplica(const UUID& client_id, } if (!metadata.IsValid()) { + PublishKvRemoved(key, metadata, tenant_id); accessor.Erase(); } return {}; @@ -4495,6 +4501,7 @@ auto MasterService::Remove(const std::string& key, const std::string& tenant_id, return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK); } + PublishKvRemoved(key, metadata, tenant_id); auto& tenant_state = accessor.GetTenantState(); accessor.Erase(); return {}; @@ -7344,6 +7351,8 @@ void MasterService::BatchEvict(double evict_ratio_target, result.freed_bytes += freed; if (freed > 0) { result.evicted_objects++; + PublishKvRemovedAfterEvict(member_key, freed, "cpu", + member_metadata, tenant_id); } if (member_key != key && !member_metadata.IsValid()) { EraseMetadata(tenant_state, member_it, tenant_id, @@ -7501,6 +7510,10 @@ void MasterService::BatchEvict(double evict_ratio_target, deferred_replicas, /*allow_soft_pinned=*/false); total_freed_size += evict_result.freed_bytes; + if (!it->second.IsGrouped()) { + PublishKvRemovedAfterEvict(c.key, evict_result.freed_bytes, + "cpu", it->second, c.tenant_id); + } if (!it->second.IsValid()) { EraseMetadata(tenant_state, it, c.tenant_id, QuotaEraseMode::kFull, &shard); @@ -7561,6 +7574,11 @@ void MasterService::BatchEvict(double evict_ratio_target, shard, tenant_state, deferred_replicas, /*allow_soft_pinned=*/false); total_freed_size += evict_result.freed_bytes; + if (!it->second.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } if (!it->second.IsValid()) { it = EraseMetadata( tenant_state, it, tenant_it->first, @@ -7621,6 +7639,11 @@ void MasterService::BatchEvict(double evict_ratio_target, shard, tenant_state, deferred_replicas, /*allow_soft_pinned=*/true); total_freed_size += evict_result.freed_bytes; + if (!it->second.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } if (!it->second.IsValid()) { it = EraseMetadata( tenant_state, it, tenant_it->first, @@ -7780,6 +7803,8 @@ void MasterService::NoFBatchEvict(double evict_ratio_target, total_freed_size += metadata.size * erased; shard_evicted_count++; + PublishKvRemovedAfterEvict(it->first, metadata.size * erased, + "disk", metadata, tenant_it->first); if (!metadata.IsValid()) { it = EraseMetadata(tenant_state, it, tenant_it->first, QuotaEraseMode::kFull, &shard); @@ -9509,6 +9534,108 @@ MasterService::MetadataSerializer::DeserializeDiscardedReplicas( return {}; } +KvEventConfig MasterService::BuildKvEventConfig( + const MasterServiceConfig& config) { + KvEventConfig kv_config; + kv_config.enabled = config.enable_kv_events; + kv_config.bind_endpoint = config.kv_events_bind_endpoint; + kv_config.model_name = config.kv_events_model_name; + kv_config.backend_id = config.kv_events_backend_id; + kv_config.tenant_id = config.kv_events_tenant_id; + kv_config.additional_salt = config.kv_events_additional_salt; + kv_config.lora_name = config.kv_events_lora_name; + kv_config.block_size = config.kv_events_block_size; + kv_config.dp_rank = config.kv_events_dp_rank; + kv_config.emit_legacy_compat_fields = config.kv_events_emit_legacy_compat; + kv_config.emit_object_key = config.kv_events_emit_object_key; + kv_config.queue_capacity = config.kv_events_queue_capacity; + return kv_config; +} + +std::string MasterService::MediumForReplicaType(ReplicaType replica_type) { + switch (replica_type) { + case ReplicaType::MEMORY: + return "cpu"; + case ReplicaType::DISK: + case ReplicaType::LOCAL_DISK: + case ReplicaType::NOF_SSD: + return "disk"; + case ReplicaType::ALL: + default: + return "cpu"; + } +} + +std::string MasterService::MediumForMetadata(const ObjectMetadata& metadata) { + if (metadata.HasMemReplica()) { + return "cpu"; + } + if (metadata.HasReplica(&Replica::fn_is_nof_replica) || + metadata.HasReplica(&Replica::fn_is_disk_replica) || + metadata.HasReplica(&Replica::fn_is_local_disk_replica)) { + return "disk"; + } + return "cpu"; +} + +void MasterService::PublishKvStored(const std::string& key, + ReplicaType replica_type, + const ObjectMetadata& metadata, + const std::string& tenant_id) { + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + std::string medium = MediumForReplicaType(replica_type); + if (replica_type == ReplicaType::ALL) { + medium = MediumForMetadata(metadata); + } + kv_event_publisher_->PublishStored(key, medium, tenant_id, + metadata.group_id); +} + +void MasterService::PublishKvRemoved(const std::string& key, + const std::string& medium, + const std::string& tenant_id, + const std::string& group_id) { + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + kv_event_publisher_->PublishRemoved(key, medium, tenant_id, group_id); +} + +void MasterService::PublishKvRemoved(const std::string& key, + const ObjectMetadata& metadata, + const std::string& tenant_id) { + PublishKvRemoved(key, MediumForMetadata(metadata), tenant_id, + metadata.group_id); +} + +void MasterService::PublishKvRemovedAfterEvict(const std::string& key, + uint64_t freed_bytes, + const std::string& medium, + const ObjectMetadata& metadata, + const std::string& tenant_id) { + (void)freed_bytes; + (void)medium; + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + if (!metadata.IsValid()) { + PublishKvRemoved(key, metadata, tenant_id); + } +} + +bool MasterService::KvEventsEnabled() const { + return kv_event_publisher_ && kv_event_publisher_->enabled(); +} + +KvEventPublisher::Stats MasterService::GetKvEventStats() const { + if (!kv_event_publisher_) { + return {}; + } + return kv_event_publisher_->GetStats(); +} + void MasterService::setHttpMetadataServer(HttpMetadataServer* server) { http_metadata_server_ = server; if (server) { diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 5d2ae9efad..d5de362f54 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -1258,6 +1258,14 @@ WrappedMasterService::QuerySegmentStatusById(const UUID& segment_id) { return master_service_.QuerySegmentStatusById(segment_id); } +bool WrappedMasterService::KvEventsEnabled() const { + return master_service_.KvEventsEnabled(); +} + +KvEventPublisher::Stats WrappedMasterService::GetKvEventStats() const { + return master_service_.GetKvEventStats(); +} + void RegisterRpcService( coro_rpc::coro_rpc_server& server, mooncake::WrappedMasterService& wrapped_master_service) { diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index b661521adf..51f3983ff7 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -38,6 +38,21 @@ add_store_test(runtime_accelerator_test runtime_accelerator_test.cpp) add_store_test(allocation_strategy_test allocation_strategy_test.cpp) add_store_test(eviction_strategy_test eviction_strategy_test.cpp) add_store_test(deadline_scheduler_test deadline_scheduler_test.cpp) +add_store_test(kv_event_publisher_test kv_event_publisher_test.cpp) +if(ENABLE_KV_EVENTS) + find_library(KV_EVENT_TEST_ZMQ_LIBRARY NAMES zmq libzmq PATHS /usr/lib + /usr/local/lib /usr/lib64) + find_path(KV_EVENT_TEST_ZMQ_INCLUDE_DIR NAMES zmq.h PATHS /usr/include + /usr/local/include) + target_compile_definitions(kv_event_publisher_test PRIVATE + MOONCAKE_ENABLE_KV_EVENTS=1) + if(KV_EVENT_TEST_ZMQ_LIBRARY AND KV_EVENT_TEST_ZMQ_INCLUDE_DIR) + target_include_directories(kv_event_publisher_test + PRIVATE ${KV_EVENT_TEST_ZMQ_INCLUDE_DIR}) + target_link_libraries(kv_event_publisher_test + PRIVATE ${KV_EVENT_TEST_ZMQ_LIBRARY}) + endif() +endif() add_store_test(master_service_test master_service_test.cpp) add_store_test(master_service_tenant_quota_test master_service_tenant_quota_test.cpp) diff --git a/mooncake-store/tests/kv_event_publisher_test.cpp b/mooncake-store/tests/kv_event_publisher_test.cpp new file mode 100644 index 0000000000..fec8292004 --- /dev/null +++ b/mooncake-store/tests/kv_event_publisher_test.cpp @@ -0,0 +1,218 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "kv_event/key_util.h" +#include "kv_event/kv_event_publisher.h" + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS +#include +#include +#endif + +namespace mooncake { +namespace { + +TEST(KvEventKeyUtilTest, ParseSeqHashFromObjectKey) { + EXPECT_EQ(ParseSeqHashFromObjectKey("12345"), 12345u); + EXPECT_EQ(ParseSeqHashFromObjectKey("0x2a"), 42u); + EXPECT_EQ(ParseSeqHashFromObjectKey("0XFF"), 255u); + EXPECT_FALSE(ParseSeqHashFromObjectKey("").has_value()); + EXPECT_FALSE(ParseSeqHashFromObjectKey("not-a-hash").has_value()); + EXPECT_FALSE(ParseSeqHashFromObjectKey("123abc").has_value()); + EXPECT_EQ(KvEventPublisher::ParseSeqHashFromObjectKey("99"), 99u); +} + +TEST(KvEventPublisherTest, DisabledPublisherIsNoop) { + KvEventConfig config; + config.enabled = false; + KvEventPublisher publisher(config); + EXPECT_FALSE(publisher.enabled()); + publisher.PublishStored("42", "cpu"); + publisher.PublishRemoved("42", "cpu"); + const auto stats = publisher.GetStats(); + EXPECT_EQ(stats.published_events, 0u); + EXPECT_EQ(stats.dropped_events, 0u); +} + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +namespace { + +std::string MakeIpcEndpoint() { + return "ipc:///tmp/kv_event_test_" + std::to_string(getpid()) + "_" + + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count()); +} + +bool ReceiveZmqMultipart(void* socket, std::vector& frames) { + frames.clear(); + while (true) { + zmq_msg_t msg; + if (zmq_msg_init(&msg) != 0) { + return false; + } + const int rc = zmq_msg_recv(&msg, socket, 0); + if (rc < 0) { + zmq_msg_close(&msg); + return false; + } + const char* data = static_cast(zmq_msg_data(&msg)); + frames.emplace_back(data, data + zmq_msg_size(&msg)); + const int more = zmq_msg_more(&msg); + zmq_msg_close(&msg); + if (!more) { + break; + } + } + return true; +} + +} // namespace + +TEST(KvEventPublisherTest, PublishesSglangObjectKeyOverZmq) { + const std::string endpoint = MakeIpcEndpoint(); + const std::string object_key = + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855_0_k"; + const std::string group_id = + "sglang-hicache:" + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + + KvEventConfig config; + config.enabled = true; + config.bind_endpoint = endpoint; + config.backend_id = "mooncake-test"; + config.emit_object_key = true; + config.emit_legacy_compat_fields = true; + config.queue_capacity = 64; + KvEventPublisher publisher(config); + ASSERT_TRUE(publisher.enabled()); + + void* ctx = zmq_ctx_new(); + ASSERT_NE(ctx, nullptr); + void* sub = zmq_socket(ctx, ZMQ_SUB); + ASSERT_NE(sub, nullptr); + ASSERT_EQ(zmq_connect(sub, endpoint.c_str()), 0); + ASSERT_EQ(zmq_setsockopt(sub, ZMQ_SUBSCRIBE, "", 0), 0); + + // Allow SUB connect before first publish propagates. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + publisher.PublishStored(object_key, "cpu", "tenant-a", group_id); + + std::vector frames; + ASSERT_TRUE(ReceiveZmqMultipart(sub, frames)) << zmq_strerror(zmq_errno()); + ASSERT_EQ(frames.size(), 3u); + EXPECT_TRUE(frames[0].empty()); + ASSERT_EQ(frames[1].size(), sizeof(uint64_t)); + + const auto object_handle = + msgpack::unpack(frames[2].data(), frames[2].size()); + const auto& root = object_handle.get(); + ASSERT_EQ(root.type, msgpack::type::ARRAY); + ASSERT_EQ(root.via.array.size, 3u); + + const auto& events = root.via.array.ptr[1]; + ASSERT_EQ(events.type, msgpack::type::ARRAY); + ASSERT_EQ(events.via.array.size, 1u); + + const auto& event = events.via.array.ptr[0]; + ASSERT_EQ(event.type, msgpack::type::MAP); + + bool has_object_key = false; + bool has_group_id = false; + bool has_empty_seq_hashes = false; + std::string event_type; + std::string backend_id; + std::string tenant_id; + for (uint32_t i = 0; i < event.via.map.size; ++i) { + const auto& key = event.via.map.ptr[i].key; + const auto& val = event.via.map.ptr[i].val; + ASSERT_EQ(key.type, msgpack::type::STR); + const std::string field(key.via.str.ptr, key.via.str.size); + if (field == "object_key") { + ASSERT_EQ(val.type, msgpack::type::STR); + EXPECT_EQ(std::string(val.via.str.ptr, val.via.str.size), + object_key); + has_object_key = true; + } else if (field == "group_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + EXPECT_EQ(std::string(val.via.str.ptr, val.via.str.size), group_id); + has_group_id = true; + } else if (field == "seq_hashes") { + ASSERT_EQ(val.type, msgpack::type::ARRAY); + EXPECT_EQ(val.via.array.size, 0u); + has_empty_seq_hashes = true; + } else if (field == "event_type") { + ASSERT_EQ(val.type, msgpack::type::STR); + event_type = std::string(val.via.str.ptr, val.via.str.size); + } else if (field == "backend_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + backend_id = std::string(val.via.str.ptr, val.via.str.size); + } else if (field == "tenant_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + tenant_id = std::string(val.via.str.ptr, val.via.str.size); + } + } + + EXPECT_EQ(event_type, "stored"); + EXPECT_EQ(backend_id, "mooncake-test"); + EXPECT_EQ(tenant_id, "tenant-a"); + EXPECT_TRUE(has_object_key); + EXPECT_TRUE(has_group_id); + EXPECT_TRUE(has_empty_seq_hashes); + + // Wait for async worker to finish publishing. + for (int i = 0; i < 50; ++i) { + const auto stats = publisher.GetStats(); + if (stats.published_events == 1 && stats.published_batches == 1) { + EXPECT_EQ(stats.dropped_events, 0u); + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + const auto stats = publisher.GetStats(); + EXPECT_EQ(stats.published_events, 1u); + EXPECT_EQ(stats.published_batches, 1u); + EXPECT_EQ(stats.dropped_events, 0u); + + zmq_close(sub); + zmq_ctx_destroy(ctx); +} + +TEST(KvEventPublisherTest, DropsOldestWhenQueueFull) { + const std::string endpoint = MakeIpcEndpoint(); + + KvEventConfig config; + config.enabled = true; + config.bind_endpoint = endpoint; + config.backend_id = "mooncake-test"; + config.emit_object_key = true; + config.queue_capacity = 2; + KvEventPublisher publisher(config); + ASSERT_TRUE(publisher.enabled()); + + for (int i = 0; i < 100; ++i) { + publisher.PublishStored(std::to_string(i), "cpu"); + } + + for (int i = 0; i < 50; ++i) { + const auto stats = publisher.GetStats(); + if (stats.dropped_events >= 1) { + EXPECT_GE(stats.dropped_events, 1u); + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + FAIL() << "expected dropped_events after queue overflow"; +} + +#endif // MOONCAKE_ENABLE_KV_EVENTS + +} // namespace +} // namespace mooncake From 713edc524df17aee290418a96b9648f2a6d4a50f Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Wed, 8 Jul 2026 15:12:58 +0800 Subject: [PATCH 34/57] [TENT] RailMonitor: prefer same-name device for cross-NUMA rail mapping (#2758/#2467) (#2790) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updateBestMapping() maps a local NIC to a remote NIC per (local_numa, remote_numa) pair. For the same-NUMA case it uses direct_rails_, which loadDefault() builds via same-name matching (mlx5_5 -> mlx5_5). But for the cross-NUMA case it used positional assignment, remote_devices[remote_numa][i % remote_cnt], ignoring device names. On a multi-bond dual-NUMA RoCEv2 fabric where the two nodes disagree on which NUMA a same-named NIC sits in (e.g. an overlay NIC), this maps a local NIC to an unrelated remote NIC on a different physical/overlay network. The QP then never reaches RTR -> 'transport retry counter exceeded' (#2467) / cross-node modify-to-RTR EINVAL(22) (#2758). Fix: in the cross-NUMA branch, prefer a same-name remote device (mirroring loadDefault()'s Priority-1 matching) before falling back to positional assignment. Same-NUMA behavior is unchanged; positional fallback still applies when no same-name remote device exists in the target NUMA domain. Add a cross-NUMA same-name unit test (asymmetric NUMA layout as in #2467). Verified: the changed translation unit compiles cleanly in-tree; the full tent_rail_monitor_test binary could not be linked in my environment due to an unrelated GDS/cuFile build dependency (cufile.h absent), so the gtest was not run locally. Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 --- .../tent/src/transport/rdma/rail_monitor.cpp | 28 ++++++++- .../tent/tests/rail_monitor_test.cpp | 63 +++++++++++++++++++ 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp index 438b211dcc..21e6f904f5 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp @@ -323,10 +323,32 @@ void RailMonitor::updateBestMapping() { for (size_t i = 0; i < local_cnt; i++) { int local_nic = local_devices[local_numa][i]; int remote_nic = -1; - if (local_numa == remote_numa) + if (local_numa == remote_numa) { remote_nic = direct_rails_[local_nic]; - else - remote_nic = remote_devices[remote_numa][i % remote_cnt]; + } else { + // Cross-NUMA: prefer a same-name remote device (e.g. + // mlx5_5 -> mlx5_5) before falling back to positional + // assignment. loadDefault() already builds direct_rails_ + // via same-name matching (Priority 1); mirroring it here + // avoids mapping a local NIC to an unrelated remote NIC on + // a different physical/overlay network, which fails QP + // modify-to-RTR with "transport retry counter exceeded" on + // multi-bond dual-NUMA RoCEv2 fabrics (issues #2758/#2467). + auto local_entry = local_->getNicEntry(local_nic); + if (local_entry) { + for (int cand : remote_devices[remote_numa]) { + auto cand_entry = remote_->getNicEntry(cand); + if (cand_entry && + cand_entry->name == local_entry->name) { + remote_nic = cand; + break; + } + } + } + if (remote_nic < 0) + remote_nic = + remote_devices[remote_numa][i % remote_cnt]; + } if (!available(local_nic, remote_nic)) { bool found = false; for (int cand : remote_devices[remote_numa]) { diff --git a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp index 5963ed008d..fc05c7a411 100644 --- a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp @@ -104,6 +104,69 @@ TEST(RailMonitorConfigTest, CustomJsonOverridesAutomaticPeerMapping) { EXPECT_FALSE(rail.available(/*local_nic=*/0, /*remote_nic=*/0)); } +// Build a 2-NIC topology (mlx5_a, mlx5_b) with per-NIC NUMA nodes, so the two +// sides can disagree on which NUMA a same-named NIC sits in — the asymmetric +// (overlay) situation from #2467. +static std::shared_ptr makeNamedNumaTopology(const std::string& n0, + int numa0, + const std::string& n1, + int numa1) { + auto json_str = + R"({ + "nics": [ + {"name": ")" + + n0 + R"(", "type": 0, "numa_node": )" + std::to_string(numa0) + R"(}, + {"name": ")" + + n1 + R"(", "type": 0, "numa_node": )" + std::to_string(numa1) + R"(} + ], + "mems": [{ + "name": "host0", + "type": 0, + "numa_node": 0, + "device_list": {"rank0": [0, 1]} + }] + })"; + auto topo = std::make_shared(); + auto status = topo->parse(json_str); + if (!status.ok()) { + ADD_FAILURE() << "Topology::parse failed: " << status.ToString(); + } + return topo; +} + +// --------------------------------------------------------------------------- +// Cross-NUMA mapping must prefer a same-name remote device over a positional +// (i % remote_cnt) pick, so a local NIC is not routed to an unrelated remote +// NIC on a different physical/overlay network (issues #2758/#2467). +// +// Setup (asymmetric NUMA, as in #2467's overlay case): +// local : mlx5_x @ NUMA 0 (idx0), mlx5_y @ NUMA 1 (idx1) +// remote: mlx5_y @ NUMA 0 (idx0), mlx5_x @ NUMA 1 (idx1) +// Local mlx5_y sits in NUMA 1; its same-name remote mlx5_y sits in NUMA 0. +// Querying local mlx5_y (idx1) for the remote NUMA-0 domain is cross-NUMA and +// must pick the same-name remote mlx5_y (remote idx0). The positional bug would +// instead pick remote_devices[NUMA0][i]. With only one device in that domain +// they coincide, so we make the discriminating assertion below. +// --------------------------------------------------------------------------- + +TEST(RailMonitorCrossNumaTest, CrossNumaPrefersSameNameDevice) { + // local NUMA-1 domain has one NIC: mlx5_y (idx1). + auto local = makeNamedNumaTopology("mlx5_x", 0, "mlx5_y", 1); + // remote NUMA-0 domain: mlx5_y (idx0); remote NUMA-1 domain: mlx5_x (idx1). + auto remote = makeNamedNumaTopology("mlx5_y", 0, "mlx5_x", 1); + RailMonitor rail; + ASSERT_TRUE(rail.load(local.get(), remote.get()).ok()); + ASSERT_TRUE(rail.ready()); + + // local mlx5_y (idx1, NUMA 1) reaching the remote NUMA-0 domain: the only + // same-name device is remote mlx5_y at idx0. Must map there. + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/1, /*remote_numa=*/0), 0); + + // local mlx5_x (idx0, NUMA 0) reaching remote NUMA-1 domain: same-name + // remote mlx5_x is at idx1. Must map there, not positionally to idx0. + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/0, /*remote_numa=*/1), 1); +} + // --------------------------------------------------------------------------- // markRecovered resets error_count so failures start accumulating fresh // --------------------------------------------------------------------------- From e07774fb35650ac76bed0b995e039664376ea3a9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:27:01 +0800 Subject: [PATCH 35/57] [Build] Bump golang.org/x/net in /mooncake-p2p-store/src/p2pstore (#2708) Bumps [golang.org/x/net](https://github.com/golang/net) from 0.48.0 to 0.55.0. - [Commits](https://github.com/golang/net/compare/v0.48.0...v0.55.0) --- updated-dependencies: - dependency-name: golang.org/x/net dependency-version: 0.55.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- mooncake-p2p-store/src/p2pstore/go.mod | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/mooncake-p2p-store/src/p2pstore/go.mod b/mooncake-p2p-store/src/p2pstore/go.mod index 5a60aa938a..5929721056 100644 --- a/mooncake-p2p-store/src/p2pstore/go.mod +++ b/mooncake-p2p-store/src/p2pstore/go.mod @@ -1,8 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-p2p-store/src/p2pstore -go 1.24.0 - -toolchain go1.24.1 +go 1.25.0 require go.etcd.io/etcd/client/v3 v3.5.15 @@ -16,9 +14,9 @@ require ( go.uber.org/atomic v1.7.0 // indirect go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.17.0 // indirect - golang.org/x/net v0.48.0 // indirect - golang.org/x/sys v0.39.0 // indirect - golang.org/x/text v0.32.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect google.golang.org/grpc v1.79.3 // indirect From f1c762a714d5bb13cc3e4ec1da7404cb5a5b6e29 Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Wed, 8 Jul 2026 16:20:57 +0800 Subject: [PATCH 36/57] [Store] fix: reject empty keys in HTTP metadata server (#2770) --------- Co-authored-by: chatgpt-codex-connector[bot] <199175422+chatgpt-codex-connector[bot]@users.noreply.github.com> --- .../mooncake/http_metadata_server.py | 5 +- .../tests/test_http_metadata_server.py | 82 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 mooncake-wheel/tests/test_http_metadata_server.py diff --git a/mooncake-wheel/mooncake/http_metadata_server.py b/mooncake-wheel/mooncake/http_metadata_server.py index cf3cf21d43..b37cfb4006 100644 --- a/mooncake-wheel/mooncake/http_metadata_server.py +++ b/mooncake-wheel/mooncake/http_metadata_server.py @@ -61,7 +61,10 @@ def _setup_routes(self): async def _handle_metadata(self, request: web.Request): """Handle metadata requests.""" - key = request.query.get('key', '') + key = request.query.get('key', '').strip() + if not key: + return web.Response(text='metadata key is required', status=400, + content_type='application/json') if request.method == 'GET': return await self._handle_get(key) diff --git a/mooncake-wheel/tests/test_http_metadata_server.py b/mooncake-wheel/tests/test_http_metadata_server.py new file mode 100644 index 0000000000..a494d5f203 --- /dev/null +++ b/mooncake-wheel/tests/test_http_metadata_server.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from mooncake.http_metadata_server import KVBootstrapServer + + +class FakeRequest: + def __init__(self, method, key=None, body=b""): + self.method = method + self.query = {} if key is None else {"key": key} + self.body = body + + async def read(self): + return self.body + + +class HttpMetadataServerTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.server = KVBootstrapServer(port=0) + + async def test_missing_metadata_key_is_rejected_for_all_methods(self): + for method in ("GET", "PUT", "DELETE"): + with self.subTest(method=method): + response = await self.server._handle_metadata( + FakeRequest(method, body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn("", self.server.store) + + async def test_empty_metadata_key_is_rejected(self): + response = await self.server._handle_metadata( + FakeRequest("PUT", key="", body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn("", self.server.store) + + async def test_blank_metadata_key_is_rejected(self): + response = await self.server._handle_metadata( + FakeRequest("PUT", key=" ", body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn(" ", self.server.store) + + async def test_metadata_key_is_stripped_before_operations(self): + put_response = await self.server._handle_metadata( + FakeRequest("PUT", key=" valid ", body=b"value") + ) + get_response = await self.server._handle_metadata( + FakeRequest("GET", key=" valid ") + ) + + self.assertEqual(put_response.status, 200) + self.assertEqual(get_response.status, 200) + self.assertEqual(get_response.body, b"value") + self.assertIn("valid", self.server.store) + self.assertNotIn(" valid ", self.server.store) + + async def test_valid_metadata_key_still_round_trips(self): + put_response = await self.server._handle_metadata( + FakeRequest("PUT", key="valid", body=b"value") + ) + get_response = await self.server._handle_metadata( + FakeRequest("GET", key="valid") + ) + + self.assertEqual(put_response.status, 200) + self.assertEqual(get_response.status, 200) + self.assertEqual(get_response.body, b"value") + + +if __name__ == "__main__": + unittest.main() From df84a553efd5572880109f1ecabd60ad1549380b Mon Sep 17 00:00:00 2001 From: lujh <101535776+LujhCoconut@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:43:21 +0800 Subject: [PATCH 37/57] [wip] docs: add vLLM V1 MooncakeStore KV cache sharing benchmark (#2773) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../examples/vllm-integration/index.md | 2 +- .../vllm-mooncakestoreconnector.md | 5 +++ .../performance/vllm-v1-kvcache-sharing.md | 45 +++++++++++++++++++ docs/source/performance/vllm/index.md | 2 + 4 files changed, 53 insertions(+), 1 deletion(-) create mode 100644 docs/source/performance/vllm-v1-kvcache-sharing.md diff --git a/docs/source/getting_started/examples/vllm-integration/index.md b/docs/source/getting_started/examples/vllm-integration/index.md index d6107e502e..f09f4c65b6 100644 --- a/docs/source/getting_started/examples/vllm-integration/index.md +++ b/docs/source/getting_started/examples/vllm-integration/index.md @@ -5,7 +5,7 @@ Mooncake integrates with vLLM to accelerate large language model serving through high-performance KV cache transfer and shared storage. The integration supports two primary scenarios: - **Disaggregated Prefill-Decode Serving**: Seamlessly split prefill and decode across nodes using `MooncakeConnector`, with RDMA-powered cross-node KV cache transfer achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE). Transfer overhead is negligible — for 32K-token prompts (4.50 GB of KV data), transfer takes only **31.65 ms**, accounting for just **4.2%** of total TTFT. -- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime. +- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime. Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x** (1P1D, 12GPUs), and scales to **60 GPUs** with >95% cache hit rate as shown in this [webpage](../../../../../docs/source/performance/vllm-v1-kvcache-sharing.md). | Scenario | Guide | vLLM Backend | |----------|-------|-------------| diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md index bb0a5f52fc..25dff8fc4b 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md @@ -137,3 +137,8 @@ python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_ > ``` > > Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits. + + +### 4. Performance + +Please refer to this [webpage](../../../../../docs/source/performance/vllm-v1-kvcache-sharing.md). \ No newline at end of file diff --git a/docs/source/performance/vllm-v1-kvcache-sharing.md b/docs/source/performance/vllm-v1-kvcache-sharing.md new file mode 100644 index 0000000000..8b63999ea9 --- /dev/null +++ b/docs/source/performance/vllm-v1-kvcache-sharing.md @@ -0,0 +1,45 @@ +# Benchmark performance +Mooncake leverages the `MooncakeStoreConnector` in vLLM V1 to enable a distributed KV cache pool, supporting cross-instance sharing and reuse of KV caches. Furthermore, vLLM's `MultiConnector` can be configured to orchestrate both the `MooncakeConnector` (for peer-to-peer KV transfer) and the `MooncakeStoreConnector` (for the shared pool), enabling prefill-decode (PD) disaggregation. + +![Overall Performance](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/hero_vllm_mooncake.svg) + +We thank the vLLM team for conducting the performance evaluation. The detailed results are presented below. + +> The original blog is available at https://vllm.ai/blog/2026-05-06-mooncake-store. + + +## Speeding up real agentic traces + +Setup: Kimi-2.5 NVFP4 model on GB200 nodes with PD disaggregation + +In this experiment, the model was deployed with a 1P1D configuration across 12 GPUs in total. + +![Throughput and Latency Comparison on Agentic Traces](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/pd_compare_mooncake_vs_nixl.png) + +The distributed KV cache pool improves vLLM throughput by 3.8x and reduces P50 TTFT and E2E latency by 46x and 8.6x, respectively. These gains are driven by a dramatic increase in cache hit rate: from 1.7%, where only the system prompt is cached, to 92.2%, where nearly the entire prefix is cached. + +## Scaling out to multiple nodes + +Experiment settings: + +* 20K common tokens (system instructions) +* 10K tokens first input +* 2,048 tokens per-turn input length +* 900 output tokens +* 30 turns total +* Number of sessions scaled with number of GPUs: 75 → 150 → 225 → 300 → 375 +* Parameters were chosen to roughly align with the original Codex workload and keep the total output/input ratio ~1.3% + +![Scaling Performance across Multiple Nodes](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/pd_scaling.png) + +To stress-test the datapath under cross-node traffic, we used round-robin routing. As a result, requests could be scheduled on different nodes across turns and often needed to fetch KV caches from a previous node. + +Without a distributed KV cache pool, this routing pattern would cause massive cache misses and severe throughput degradation. With Mooncake Store, vLLM consistently achieves a cache hit rate above 95%, and the system scales nearly linearly to 60 GPUs. + +This result shows that the distributed KV cache pool substantially improves cache hit rate while maintaining an efficient datapath as the cluster grows. + + +## Benchmark Scripts + +The benchmark scripts are provided in the artifact repository [here](https://github.com/ivanium/vllm/tree/feat/mooncake-store-int/scripts/mooncake/artifacts). + diff --git a/docs/source/performance/vllm/index.md b/docs/source/performance/vllm/index.md index b32561c7dc..eb9780e5c1 100644 --- a/docs/source/performance/vllm/index.md +++ b/docs/source/performance/vllm/index.md @@ -4,6 +4,7 @@ Benchmarks evaluating Mooncake's integration with vLLM across different backends | Document | Backend | Key Findings | |----------|---------|---------------| +| [vLLM V1 + MooncakeStoreConnector](../vllm-v1-kvcache-sharing) | vLLM V1 | Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x**, and scales to **60 GPUs** with >95% cache hit rate | | [vLLM V1 + MooncakeConnector](../vllm-v1-support-benchmark) | vLLM V1 | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens | | [vLLM V1 + MooncakeStore vs Redis](../vllm-benchmark-results-v1) | vLLM V1 | MooncakeStore RDMA consistently outperforms Redis across all XpYd topologies — e.g., **~32% lower** mean TTFT in 2P2D tp=2 | | [vLLM V0 + MooncakeConnector (Legacy)](../vllm-benchmark-results-v0.2) | vLLM V0 | TP=4 reduces TTFT by ~80% vs TP=1; RDMA provides significant latency advantage over TCP across varying QPS and input lengths | @@ -12,6 +13,7 @@ Benchmarks evaluating Mooncake's integration with vLLM across different backends :maxdepth: 1 :hidden: +../vllm-v1-kvcache-sharing ../vllm-v1-support-benchmark ../vllm-benchmark-results-v1 ../vllm-benchmark-results-v0.2 From c9896684fbd7b85ca207c643056a645ab6be3bad Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:57:24 +0800 Subject: [PATCH 38/57] [Build] Bump golang.org/x/crypto (#2789) Bumps [golang.org/x/crypto](https://github.com/golang/crypto) from 0.51.0 to 0.52.0. - [Commits](https://github.com/golang/crypto/compare/v0.51.0...v0.52.0) --- updated-dependencies: - dependency-name: golang.org/x/crypto dependency-version: 0.52.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- mooncake-transfer-engine/example/http-metadata-server/go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mooncake-transfer-engine/example/http-metadata-server/go.mod b/mooncake-transfer-engine/example/http-metadata-server/go.mod index 2646bb118a..3641bc48c0 100644 --- a/mooncake-transfer-engine/example/http-metadata-server/go.mod +++ b/mooncake-transfer-engine/example/http-metadata-server/go.mod @@ -25,7 +25,7 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect golang.org/x/arch v0.8.0 // indirect - golang.org/x/crypto v0.51.0 // indirect + golang.org/x/crypto v0.52.0 // indirect golang.org/x/net v0.55.0 // indirect golang.org/x/sys v0.45.0 // indirect golang.org/x/text v0.37.0 // indirect From 866ed32f6f8c351b900b805f9c15b00bdf0f5d07 Mon Sep 17 00:00:00 2001 From: tancz <544463199@qq.com> Date: Thu, 9 Jul 2026 11:27:32 +0800 Subject: [PATCH 39/57] [Store] Fix --host parameter to support ip:port format for TransferEngine data plane port (#2784) Previously, passing --host=ip:port would cause "IPC server stopped" because the port-containing string was forwarded to coro_rpc_server without stripping the port. Now getHostNameWithoutPort() is used to extract the bare IP for coro_rpc_server binding, while the port flows through to TransferEngine as intended. Signed-off-by: tan changzhi <544463199@qq.com> --- .../mooncake-store-deployment-guide.md | 4 +- docs/source/deployment/ssd-offload.md | 2 +- mooncake-store/src/real_client_main.cpp | 6 ++- mooncake-store/tests/CMakeLists.txt | 1 + mooncake-store/tests/host_port_fix_test.cpp | 48 +++++++++++++++++++ 5 files changed, 56 insertions(+), 5 deletions(-) create mode 100644 mooncake-store/tests/host_port_fix_test.cpp diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 9798cf3d23..c012e714d3 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -735,8 +735,8 @@ mooncake_client \ | Flag | Default | Description | |------|---------|-------------| -| `--host` | `0.0.0.0` | Client service bind host | -| `--port` | `50052` | Client service listen port | +| `--host` | `0.0.0.0` | Client service bind host. Accepts `ip:port` to specify the data plane port for TransferEngine | +| `--port` | `50052` | Client RPC listen port (dummy↔real client control plane) | | `--global_segment_size` | `4 GB` | Global segment size contributed by the client | | `--master_server_address` | `127.0.0.1:50051` | Master service address | | `--metadata_server` | `http://127.0.0.1:8080/metadata` | Transfer Engine metadata service | diff --git a/docs/source/deployment/ssd-offload.md b/docs/source/deployment/ssd-offload.md index 0a5aedd08d..d1ffe86b29 100644 --- a/docs/source/deployment/ssd-offload.md +++ b/docs/source/deployment/ssd-offload.md @@ -94,7 +94,7 @@ store.setup_dummy( |------|---------|-------------| | `--metadata_server` | `http://127.0.0.1:8080/metadata` | Metadata server connection string | | `--master_server_address` | `127.0.0.1:50051` | Master address | -| `--host` | `0.0.0.0` | This machine's externally reachable IP | +| `--host` | `0.0.0.0` | This machine's externally reachable IP. Accepts `ip:port` to specify the data plane port for TransferEngine | | `--port` | `50052` | Real client RPC listening port | | `--device_names` | ` ` | NIC name(s), e.g. `eth0` or `mlx5_0` | | `--protocol` | `tcp` | Transport protocol: `tcp` or `rdma` | diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index e912e0604d..b003856662 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -3,6 +3,7 @@ #include #include "client_service.h" +#include "common.h" #include "config.h" #include "real_client.h" @@ -128,10 +129,11 @@ int main(int argc, char *argv[]) { return -1; } - coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host); + auto rpc_bind_host = getHostNameWithoutPort(FLAGS_host); + coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, rpc_bind_host); RegisterClientRpcService(server, *client_inst); - LOG(INFO) << "Starting real client service on " << FLAGS_host << ":" + LOG(INFO) << "Starting real client service on " << rpc_bind_host << ":" << FLAGS_port; return server.start(); diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 51f3983ff7..fa664ddf19 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -87,6 +87,7 @@ add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp) add_store_test(client_tcp_local_memcpy_test client_tcp_local_memcpy_test.cpp) add_store_test(pybind_client_test pybind_client_test.cpp) add_store_test(ipv6_client_test ipv6_client_test.cpp) +add_store_test(host_port_fix_test host_port_fix_test.cpp) add_store_test(client_metrics_test client_metrics_test.cpp) add_store_test(ssd_metrics_test ssd_metrics_test.cpp) add_store_test(serializer_test serializer_test.cpp) diff --git a/mooncake-store/tests/host_port_fix_test.cpp b/mooncake-store/tests/host_port_fix_test.cpp new file mode 100644 index 0000000000..6609ef9e7d --- /dev/null +++ b/mooncake-store/tests/host_port_fix_test.cpp @@ -0,0 +1,48 @@ +#include +#include +#include + +#include "common.h" +#include "utils.h" + +namespace mooncake { +namespace { + +// Regression: --host=ip:port should not break coro_rpc_server binding. +// When --host includes a port (e.g. 127.0.0.1:18007) for the TransferEngine +// data plane, the port must be stripped before the address is passed to +// coro_rpc_server, otherwise the server fails to start. + +class HostPortFixTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + FLAGS_logtostderr = true; + google::InitGoogleLogging("HostPortFixTest"); + } + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } +}; + +TEST_F(HostPortFixTest, StripsPortForRpcServer) { + auto rpc_bind_host = getHostNameWithoutPort("127.0.0.1:18007"); + EXPECT_EQ(rpc_bind_host, "127.0.0.1"); + + auto port = getFreeTcpPort(); + coro_rpc::coro_rpc_server server(/*thread_num=*/1, port, rpc_bind_host); + auto ec = server.async_start(); + EXPECT_FALSE(ec.hasResult()) << "Server should start with bare hostname"; + server.stop(); +} + +TEST_F(HostPortFixTest, BareHostPassesThrough) { + auto rpc_bind_host = getHostNameWithoutPort("127.0.0.1"); + EXPECT_EQ(rpc_bind_host, "127.0.0.1"); + + auto port = getFreeTcpPort(); + coro_rpc::coro_rpc_server server(/*thread_num=*/1, port, rpc_bind_host); + auto ec = server.async_start(); + EXPECT_FALSE(ec.hasResult()) << "Server should start with bare hostname"; + server.stop(); +} + +} // namespace +} // namespace mooncake From da11646588d8c275d992e9cfeffc44862049ac2b Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Fri, 10 Jul 2026 09:52:25 +0800 Subject: [PATCH 40/57] [TENT] Opt-in per-entry priority promotion (#2528) (#2788) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Opt-in per-entry priority promotion (#2528) Workers::promoteTimedOutRequests drains a whole priority queue but decides promotion from the HEAD entry only, then promotes every entry when the head has timed out. This over-promotes freshly enqueued, non-starving entries and, conversely, ignores timed-out entries behind a fresh head. It also promotes at most one level per tick. Factor the promotion decision into promotion_policy.h (DecidePromotionHeadOnly = today's behavior, DecidePromotionPerEntry = promote exactly the timed-out entries) so it is unit-testable without the RDMA stack, and wire an opt-in config flag: * transports/rdma/priority_promotion_per_entry (default false) keeps the historical head-only 'flush the tier' policy byte-for-byte; * when true, each pass promotes only the entries that have themselves timed out, and both MEDIUM->HIGH and LOW->MEDIUM are considered each tick so a starving LOW entry is not stalled behind an unrelated MEDIUM promotion. promotion_policy_test adds 5 deterministic cases reproducing both defects and verifying the two policies agree on the all-timed-out / empty / no-timestamp cases the head-only design targets. See issue #2528. Co-Authored-By: Claude Opus 4.8 * [TENT] per-entry promotion: add config docs + guard unsigned underflow Addresses review feedback on #2528/#2788: * docs/design/tent/qos.md: document transports/rdma/priority_promotion_per_entry (default false = head-only, true = per-entry) and the existing priority_promotion_timeout_us (per @alogfans). * promotion_policy.h: guard current_ts >= ts before the unsigned subtraction so a non-monotonic clock / race (current_ts < enqueue_ts) cannot underflow and spuriously mark an entry timed out. Co-Authored-By: Claude Opus 4.8 * fix: guard promotion decision indices * perf: keep default promotion path allocation-light --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 Co-authored-by: Yanshu <237344440@qq.com> --- docs/source/design/tent/qos.md | 14 +++ .../tent/transport/rdma/promotion_policy.h | 92 +++++++++++++++++ .../include/tent/transport/rdma/workers.h | 5 + .../tent/src/transport/rdma/workers.cpp | 83 ++++++++++------ .../tent/tests/CMakeLists.txt | 6 ++ .../tent/tests/promotion_policy_test.cpp | 99 +++++++++++++++++++ 6 files changed, 271 insertions(+), 28 deletions(-) create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h create mode 100644 mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp diff --git a/docs/source/design/tent/qos.md b/docs/source/design/tent/qos.md index 06ecfc27db..a468b4145a 100644 --- a/docs/source/design/tent/qos.md +++ b/docs/source/design/tent/qos.md @@ -62,6 +62,20 @@ This ensures that: - High-priority requests normally never wait behind lower-priority work - Low-priority requests eventually get serviced even under continuous high-priority load +**Promotion configuration**: + +| Config key | Default | Behavior | +|---|---|---| +| `transports/rdma/priority_promotion_timeout_us` | `10000` (10ms) | How long an entry may wait before it is eligible for promotion. | +| `transports/rdma/priority_promotion_per_entry` | `false` | Selects the promotion policy (see below). | + +`priority_promotion_per_entry` controls *which* entries a promotion pass moves up: + +- **`false` (default, head-only)**: a pass inspects only the queue head; if the head has timed out, the whole queue is promoted one level. This is the original, lowest-overhead "flush the tier" behavior — coarse, but it never scans the queue. +- **`true` (per-entry)**: a pass promotes exactly the entries that have themselves timed out, leaving freshly enqueued entries in place, and considers both MEDIUM→HIGH and LOW→MEDIUM in the same tick. This avoids promoting non-starving requests and avoids stalling a starving LOW entry behind an unrelated MEDIUM promotion, at the cost of scanning the drained queue. Behavior is identical to head-only for the all-timed-out / empty cases. + +The default is byte-for-byte the historical behavior; set the flag to `true` to opt into finer-grained, fairer promotion. + ### Global Slot Coordination For multi-process environments, TENT implements global time-sliced coordination using shared memory: diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h new file mode 100644 index 0000000000..a5d3a92549 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h @@ -0,0 +1,92 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Starvation-prevention promotion policy for the RDMA worker priority queues, +// factored out of Workers::promoteTimedOutRequests so the decision logic can be +// unit-tested without the full RDMA stack (issue #2528). +// +// A worker keeps three FIFO priority queues (HIGH > MEDIUM > LOW). To keep +// lower-priority requests from starving, every ~1ms a promotion pass looks at +// timed-out entries and moves them up one level. This header isolates *which* +// entries a pass decides to promote, given each entry's enqueue timestamp. + +#pragma once + +#include +#include +#include + +namespace mooncake { +namespace tent { + +// Result of a promotion decision for a single queue: the indices (into the +// drained queue order) that should move up one level. Empty == promote none. +struct PromotionDecision { + std::vector promote_indices; + bool promoted_any() const { return !promote_indices.empty(); } +}; + +// Historical policy (as implemented in Workers::promoteTimedOutRequests today): +// the pass drains the whole queue, inspects ONLY the head entry, and if the +// head has timed out it promotes EVERY entry in the queue. This is the behavior +// #2528 flags as unintended: +// * decision is head-only but applied to the whole queue -> freshly enqueued +// entries that are not starving get promoted alongside the starving head; +// * if the head has NOT timed out, later timed-out entries are not promoted. +// +// `enqueue_ts` is per drained entry in queue order (index 0 == head). A zero +// timestamp means "no timestamp" and never counts as timed out (matches the +// `enqueue_ts > 0` guard in the current code). +inline PromotionDecision DecidePromotionHeadOnly( + const std::vector& enqueue_ts, uint64_t current_ts, + uint64_t promotion_timeout_ns) { + PromotionDecision d; + if (enqueue_ts.empty()) return d; + const uint64_t head = enqueue_ts.front(); + // Guard current_ts >= head before subtracting: both are unsigned, so a + // non-monotonic clock / race where current_ts < head would otherwise + // underflow and spuriously mark the entry timed out. + const bool head_timed_out = head > 0 && current_ts >= head && + (current_ts - head) >= promotion_timeout_ns; + if (head_timed_out) { + d.promote_indices.reserve(enqueue_ts.size()); + for (size_t i = 0; i < enqueue_ts.size(); ++i) { + d.promote_indices.push_back(i); // promote ALL + } + } + return d; +} + +// Per-entry policy: promote exactly the entries that have themselves timed out, +// leaving freshly enqueued entries in place. This is the behavior #2528 +// proposes; kept here next to the historical policy so a test can contrast the +// two and a follow-up fix can switch Workers over to it. +inline PromotionDecision DecidePromotionPerEntry( + const std::vector& enqueue_ts, uint64_t current_ts, + uint64_t promotion_timeout_ns) { + PromotionDecision d; + for (size_t i = 0; i < enqueue_ts.size(); ++i) { + const uint64_t ts = enqueue_ts[i]; + // See DecidePromotionHeadOnly: guard against unsigned underflow when + // current_ts < ts. + if (ts > 0 && current_ts >= ts && + (current_ts - ts) >= promotion_timeout_ns) { + d.promote_indices.push_back(i); + } + } + return d; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h index 9b0f635daf..c18d5de130 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -215,6 +215,11 @@ class Workers { WorkerContext *worker_context_; uint64_t slice_timeout_ns_; uint64_t priority_promotion_timeout_ns_; // Timeout for priority promotion + // Opt-in (issue #2528): when true, a promotion pass promotes exactly the + // entries that have themselves timed out, instead of promoting the whole + // queue whenever only the head has timed out. Default false keeps the + // historical "flush the tier" behavior. + bool priority_promotion_per_entry_ = false; std::unique_ptr device_selector_; // File contents loaded once from workers.rail_topo_path and shared by all diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 319fec32fc..22e95b66af 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -21,6 +21,7 @@ #include #include "tent/transport/rdma/endpoint_store.h" +#include "tent/transport/rdma/promotion_policy.h" #include "tent/transport/rdma/shared_quota.h" #include "tent/common/utils/ip.h" #include "tent/common/utils/string_builder.h" @@ -142,6 +143,11 @@ Workers::Workers(RdmaTransport* transport) conf->get("transports/rdma/priority_promotion_timeout_us", 10000) * 1000ull; + // Opt-in per-entry promotion (issue #2528). Default false = historical + // head-only "flush the tier" behavior. + priority_promotion_per_entry_ = + conf->get("transports/rdma/priority_promotion_per_entry", false); + // ============================================================ // Global Slot Coordination (Multi-Process) // ============================================================ @@ -420,40 +426,61 @@ void Workers::promoteTimedOutRequests(WorkerContext& worker) { // Set next check time (1ms from now) worker.next_promotion_check_ns = current_ts + 1000000ull; - // Check MEDIUM -> HIGH promotion - std::vector promoted; - worker.queues[PRIO_MEDIUM].pop(promoted); - if (!promoted.empty()) { - auto* slice = promoted.front().first; - if (slice && slice->enqueue_ts > 0 && - (current_ts - slice->enqueue_ts) >= - priority_promotion_timeout_ns_) { - for (auto& slice_list : promoted) { - worker.queues[PRIO_HIGH].push(slice_list); + // Drain one level, promote the entries the policy selects to `to`, and put + // the rest back on `from` in their original order. Returns true if anything + // was promoted (used to preserve the historical "one level per tick" stop). + auto promote_level = [&](int from, int to) -> bool { + std::vector drained; + worker.queues[from].pop(drained); + if (drained.empty()) return false; + + if (!priority_promotion_per_entry_) { + auto* slice = drained.front().first; + const bool head_timed_out = slice && slice->enqueue_ts > 0 && + current_ts >= slice->enqueue_ts && + (current_ts - slice->enqueue_ts) >= + priority_promotion_timeout_ns_; + for (auto& slice_list : drained) { + worker.queues[head_timed_out ? to : from].push(slice_list); } - return; + return head_timed_out; } - for (auto& slice_list : promoted) { - worker.queues[PRIO_MEDIUM].push(slice_list); + + std::vector enqueue_ts; + enqueue_ts.reserve(drained.size()); + for (auto& slice_list : drained) { + auto* slice = slice_list.first; + enqueue_ts.push_back(slice ? slice->enqueue_ts : 0); } - } - // Check LOW -> MEDIUM promotion - worker.queues[PRIO_LOW].pop(promoted); - if (!promoted.empty()) { - auto* slice = promoted.front().first; - if (slice && slice->enqueue_ts > 0 && - (current_ts - slice->enqueue_ts) >= - priority_promotion_timeout_ns_) { - for (auto& slice_list : promoted) { - worker.queues[PRIO_MEDIUM].push(slice_list); - } - return; + PromotionDecision decision = DecidePromotionPerEntry( + enqueue_ts, current_ts, priority_promotion_timeout_ns_); + + if (!decision.promoted_any()) { + for (auto& slice_list : drained) + worker.queues[from].push(slice_list); + return false; } - for (auto& slice_list : promoted) { - worker.queues[PRIO_LOW].push(slice_list); + + std::vector promote(drained.size(), false); + for (size_t idx : decision.promote_indices) { + if (idx < drained.size()) promote[idx] = true; } - } + for (size_t i = 0; i < drained.size(); ++i) { + worker.queues[promote[i] ? to : from].push(drained[i]); + } + return true; + }; + + // Check MEDIUM -> HIGH promotion. Preserve the historical behavior of + // handling at most one level per tick when the head-only policy is active; + // with per-entry promotion, both levels are considered each tick so a + // starving LOW entry is not stalled behind an unrelated MEDIUM promotion. + bool promoted_medium = promote_level(PRIO_MEDIUM, PRIO_HIGH); + if (promoted_medium && !priority_promotion_per_entry_) return; + + // Check LOW -> MEDIUM promotion + promote_level(PRIO_LOW, PRIO_MEDIUM); } void Workers::asyncPollCq() { diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 868951f531..7811769a42 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -27,6 +27,12 @@ target_include_directories(admission_queue_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME admission_queue_test COMMAND admission_queue_test) +add_executable(promotion_policy_test promotion_policy_test.cpp) +target_link_libraries(promotion_policy_test PRIVATE tent_common gtest gtest_main) +target_include_directories(promotion_policy_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME promotion_policy_test COMMAND promotion_policy_test) + add_executable(tent_ip_utils_test ip_utils_test.cpp) target_link_libraries(tent_ip_utils_test PRIVATE tent_common gtest gtest_main) target_include_directories(tent_ip_utils_test diff --git a/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp new file mode 100644 index 0000000000..ab49faef73 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp @@ -0,0 +1,99 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Deterministic reproduction of the promotion issues reported in #2528. The +// historical head-only policy is contrasted with a per-entry policy on the same +// inputs so the unintended behavior is unambiguous and re-runnable (no RDMA +// stack, no timing noise). + +#include "tent/transport/rdma/promotion_policy.h" + +#include + +#include +#include + +namespace mooncake { +namespace tent { +namespace { + +constexpr uint64_t kTimeout = 10'000; // 10us promotion timeout +constexpr uint64_t kNow = 1'000'000; // fixed "now" + +// One entry that has clearly timed out. +uint64_t timedOut() { return kNow - kTimeout - 1; } +// One entry that was just enqueued and is nowhere near the timeout. +uint64_t fresh() { return kNow - 1; } + +// --- Issue #2528 point 1: head-only decision, whole-queue promotion -------- + +TEST(PromotionPolicyTest, HeadOnlyPromotesFreshEntriesWhenHeadTimedOut) { + // Queue head is starving; the two entries behind it were just enqueued. + std::vector q = {timedOut(), fresh(), fresh()}; + + auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); + + // BUG: all three are promoted, including the two fresh (non-starving) ones. + EXPECT_EQ(d.promote_indices, (std::vector{0, 1, 2})); + + // The per-entry policy promotes only the genuinely starving head. + auto fixed = DecidePromotionPerEntry(q, kNow, kTimeout); + EXPECT_EQ(fixed.promote_indices, (std::vector{0})); +} + +TEST(PromotionPolicyTest, HeadOnlyMissesTimedOutTailWhenHeadFresh) { + // Head was just enqueued; entries behind it have been starving. + std::vector q = {fresh(), timedOut(), timedOut()}; + + auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); + + // BUG: nothing is promoted even though indices 1 and 2 are starving. + EXPECT_TRUE(d.promote_indices.empty()); + + auto fixed = DecidePromotionPerEntry(q, kNow, kTimeout); + EXPECT_EQ(fixed.promote_indices, (std::vector{1, 2})); +} + +// --- Shared behavior both policies must keep ------------------------------ + +TEST(PromotionPolicyTest, EmptyQueuePromotesNothing) { + std::vector q; + EXPECT_TRUE( + DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices.empty()); + EXPECT_TRUE( + DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices.empty()); +} + +TEST(PromotionPolicyTest, ZeroTimestampNeverTimesOut) { + // enqueue_ts == 0 means "no timestamp" and must never be promoted. + std::vector q = {0, 0}; + EXPECT_TRUE( + DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices.empty()); + EXPECT_TRUE( + DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices.empty()); +} + +TEST(PromotionPolicyTest, AllTimedOutPromotesAllUnderBothPolicies) { + // When every entry is starving the two policies agree — this is the case + // the historical policy was designed around. + std::vector q = {timedOut(), timedOut(), timedOut()}; + EXPECT_EQ(DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices, + (std::vector{0, 1, 2})); + EXPECT_EQ(DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices, + (std::vector{0, 1, 2})); +} + +} // namespace +} // namespace tent +} // namespace mooncake From 6273d182157e7a904afb73bf7f9d7692ec11ec97 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Fri, 10 Jul 2026 10:14:24 +0800 Subject: [PATCH 41/57] [TENT] Expose Request.deadline_ns and policy_name to Python bindings (#2808) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Expose Request.deadline_ns and policy_name to Python bindings The C++ Request struct already carries deadline_ns (RFC #2519, #2618) and policy_name (#2640/#2759), but the Python bindings only exposed priority and transport_hint. This left deadline / policy as internal-only fields, unreachable from the Python API that upper layers (Store, SGLang, vLLM) use. Expose both as optional constructor args (defaulting to existing behavior: deadline_ns=0, policy_name=None) and as read/write properties. Fully backward compatible: callers that omit the new args are unchanged. Co-Authored-By: Claude Opus 4.8 * fix: include optional for pybind request fields --------- Co-authored-by: 彦纾 Co-authored-by: Claude Opus 4.8 Co-authored-by: Yanshu <237344440@qq.com> --- .../tent/src/python/pybind.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index bc1d3a23b7..ac7c404b67 100644 --- a/mooncake-transfer-engine/tent/src/python/pybind.cpp +++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -321,7 +322,9 @@ PYBIND11_MODULE(tent, m) { .def(py::init([](Request::OpCode opcode, uint64_t source, uint64_t target_id, uint64_t target_offset, size_t length, int priority, - TransportType transport_hint) { + TransportType transport_hint, + std::optional policy_name, + uint64_t deadline_ns) { Request r; r.opcode = opcode; r.source = U64ToPtr(source); @@ -330,12 +333,15 @@ PYBIND11_MODULE(tent, m) { r.length = length; r.priority = priority; r.transport_hint = transport_hint; + r.policy_name = std::move(policy_name); + r.deadline_ns = deadline_ns; return r; }), py::arg("opcode"), py::arg("source"), py::arg("target_id"), py::arg("target_offset"), py::arg("length"), py::arg("priority") = PRIO_HIGH, - py::arg("transport_hint") = TransportType::UNSPEC) + py::arg("transport_hint") = TransportType::UNSPEC, + py::arg("policy_name") = std::nullopt, py::arg("deadline_ns") = 0) .def_property( "opcode", [](const Request& r) { return r.opcode; }, [](Request& r, Request::OpCode op) { r.opcode = op; }) @@ -346,7 +352,9 @@ PYBIND11_MODULE(tent, m) { .def_readwrite("target_offset", &Request::target_offset) .def_readwrite("length", &Request::length) .def_readwrite("priority", &Request::priority) - .def_readwrite("transport_hint", &Request::transport_hint); + .def_readwrite("transport_hint", &Request::transport_hint) + .def_readwrite("policy_name", &Request::policy_name) + .def_readwrite("deadline_ns", &Request::deadline_ns); py::class_(m, "TransferStatus") .def(py::init<>()) From 130b959e88b2adaf131fd139b40b2817dbf83397 Mon Sep 17 00:00:00 2001 From: Yihe Liu Date: Thu, 9 Jul 2026 22:19:41 -0400 Subject: [PATCH 42/57] [Bugfix][TENT] Fix silent TPU data corruption for transfers larger than one staging chunk (#2815) Caught while testing the TPU staging path on a real TPU VM (v5p-8, libtpu 0.0.32, PJRT C API 0.83) for the first time. Until now the feature had only ever run against the mock adapter. ProxyManager stages a transfer in chunk_size (4 MiB) pieces and passes `token + chunk_offset` to the platform for every chunk after the first (proxy_manager.cpp:76). The adapter ABI only ever specified base-address classification, so `isDevicePtr(token + off)` returned false, TENT classified TPU HBM as host memory, and TpuPlatform::copy fell through to CpuPlatform::copy -- a plain memcpy of the device token. On real PJRT that memcpy does not crash: PJRT_Buffer_UnsafePointer returns an address that is host-readable but does not hold the buffer's data. The staging copy therefore produced garbage and reported COMPLETED. Chunk 0 was correct and every subsequent chunk was silently wrong, so any transfer over 4 MiB was corrupted without an error anywhere. The mock could not catch this because its "device" pointers were ordinary host memory, so the accidental memcpy happened to produce the right bytes. Fixes: - tpu_pjrt_abi.h / tpu_pjrt_shim.h: state that classification and copy entrypoints must resolve interior addresses to the registered buffer whose range contains them, that a copy may not run past a buffer's end, and that the token must never be dereferenced. - tpu_transport.cpp: require exactly one TPU-device side per staging hop and fail loudly otherwise, so a non-conforming or absent adapter can no longer degrade into a silent memcpy. - mock_tpu_pjrt_adapter.cpp: model the two properties that matter -- tokens are opaque (a poisoned read-only mapping, data lives in shadow storage) and addresses may be interior. A copy that bypasses the adapter now yields 0xDD instead of accidentally passing. - common.cmake: -DUSE_TPU=ON silently compiled zero TPU code, because all TPU sources live under tent/ which is gated on USE_TENT. This is why the bug was invisible to every build. Now a hard error. Tests: two new regression tests in tent_tpu_pjrt_shim_test and a new tent_tpu_transport_test (6 cases) that drives the staging hop the way ProxyManager does. Verified they fail against the pre-fix code (LocalStageCopiesFromInteriorDeviceOffset reports COMPLETED while delivering 0xDD) and pass after. Co-authored-by: Claude Opus 4.8 --- mooncake-common/common.cmake | 9 + .../tent/include/tent/platform/tpu_pjrt_abi.h | 32 ++- .../include/tent/platform/tpu_pjrt_shim.h | 16 +- .../tent/src/platform/tpu/README.md | 20 +- .../tent/src/transport/tpu/CMakeLists.txt | 5 +- .../tent/src/transport/tpu/tpu_transport.cpp | 30 ++- .../tent/tests/CMakeLists.txt | 14 ++ .../tent/tests/tpu/mock_tpu_pjrt_adapter.cpp | 109 ++++++++-- .../tent/tests/tpu/tpu_pjrt_shim_test.cpp | 89 +++++++- .../tent/tests/tpu/tpu_transport_test.cpp | 204 ++++++++++++++++++ 10 files changed, 486 insertions(+), 42 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index fd5292be5d..6af779a249 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -207,6 +207,15 @@ if(USE_CUDA) endif() if(USE_TPU) + # Every TPU source file lives under mooncake-transfer-engine/tent, which is + # only added when USE_TENT is ON. Without this guard -DUSE_TPU=ON configures + # and builds cleanly while compiling no TPU code at all. + if(NOT USE_TENT) + message( + FATAL_ERROR + "USE_TPU=ON requires USE_TENT=ON: all TPU support lives in TENT. Re-run cmake with -DUSE_TENT=ON." + ) + endif() add_compile_definitions(USE_TPU) message(STATUS "TPU (PJRT) staging support is enabled") endif() diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h index ceb7669785..3cb2be6fc8 100644 --- a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h @@ -24,6 +24,21 @@ // the mapping from that token to the underlying PJRT buffer; mc_tpu_pjrt_* copy // and classification calls resolve the token through the adapter's own // registry. +// +// INTERIOR POINTERS (load-bearing): TENT stages transfers through host DRAM in +// chunks (see ProxyManager, chunk_size defaults to 4 MiB), and hands the +// adapter `token + chunk_offset` for every chunk after the first. Every +// entrypoint below therefore takes an address that may point into the MIDDLE of +// a registered buffer, not only at its base. An adapter whose registry only +// matches base addresses will report `is_device_ptr(token + off) == 0`, TENT +// will classify TPU HBM as host memory, and the staging copy silently degrades +// into a memcpy from a non-data address -- corrupting every transfer larger +// than one chunk without raising an error. Adapters MUST resolve an address to +// the registered buffer whose range [base, base + size) contains it. +// +// The token is NOT required to be host-dereferenceable, and on real PJRT/TPU it +// is not: PJRT_Buffer_UnsafePointer returns an internal handle that reads as +// garbage rather than buffer contents. Never dereference it. #ifndef TENT_PLATFORM_TPU_PJRT_ABI_H_ #define TENT_PLATFORM_TPU_PJRT_ABI_H_ @@ -38,17 +53,24 @@ extern "C" { // success, non-zero on failure. Idempotent; safe to call more than once. int mc_tpu_pjrt_init(void); -// Returns 1 if `addr` is a TPU device buffer known to the adapter, else 0. +// Returns 1 if `addr` falls inside any TPU device buffer known to the adapter +// (base address or interior), else 0. int mc_tpu_pjrt_is_device_ptr(const void *addr); -// Returns the device ordinal backing `addr`, or -1 if `addr` is not a known TPU -// device buffer. +// Returns the device ordinal of the buffer containing `addr` (base address or +// interior), or -1 if `addr` is not inside a known TPU device buffer. int mc_tpu_pjrt_device_index(const void *addr); -// Synchronous device->host copy. Returns 0 on success, non-zero on failure. +// Synchronous device->host copy. `device_src` may be an interior address; the +// whole range [device_src, device_src + len) must lie within a single +// registered buffer, otherwise the adapter must fail rather than copy short. +// Returns 0 on success, non-zero on failure. int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len); -// Synchronous host->device copy. Returns 0 on success, non-zero on failure. +// Synchronous host->device copy. `device_dst` may be an interior address; the +// whole range [device_dst, device_dst + len) must lie within a single +// registered buffer, otherwise the adapter must fail rather than copy short. +// Returns 0 on success, non-zero on failure. int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len); // Number of visible TPU devices. diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h index 0d26595c7b..e2dd2d0a3d 100644 --- a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h @@ -46,18 +46,22 @@ class TpuPjrtShim { // True when the adapter library was loaded and initialized successfully. bool available() const { return available_; } - // Returns true if `addr` refers to memory owned by the TPU runtime (HBM). - // Returns false when the adapter is unavailable or the pointer is host - // memory, so a caller can safely treat "not TPU" as host memory. + // Returns true if `addr` refers to memory owned by the TPU runtime (HBM), + // including addresses interior to a registered buffer. Returns false when + // the adapter is unavailable or the pointer is host memory, so a caller can + // safely treat "not TPU" as host memory. bool isDevicePtr(const void *addr) const; - // Device ordinal backing `addr`, or -1 if `addr` is not TPU device memory. + // Device ordinal of the buffer containing `addr` (base or interior), or -1 + // if `addr` is not TPU device memory. int deviceIndex(const void *addr) const; - // Synchronous HBM -> host DMA copy of `length` bytes. + // Synchronous HBM -> host DMA copy of `length` bytes. `device_src` may be + // interior to a registered buffer; the range must not run past its end. Status copyD2H(void *host_dst, const void *device_src, size_t length) const; - // Synchronous host -> HBM DMA copy of `length` bytes. + // Synchronous host -> HBM DMA copy of `length` bytes. `device_dst` may be + // interior to a registered buffer; the range must not run past its end. Status copyH2D(void *device_dst, const void *host_src, size_t length) const; // Number of visible TPU devices (0 when the adapter is unavailable). diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/README.md b/mooncake-transfer-engine/tent/src/platform/tpu/README.md index 71f3cd4483..1cd1f068ff 100644 --- a/mooncake-transfer-engine/tent/src/platform/tpu/README.md +++ b/mooncake-transfer-engine/tent/src/platform/tpu/README.md @@ -59,10 +59,22 @@ int mc_tpu_pjrt_device_numa(int index); - **Pointer tokens:** the `const void *` "device pointer" is the stable token the serving-engine integration registers with TENT for a TPU buffer (via a `tpu:N` location). The adapter owns the mapping from that token to the - underlying PJRT buffer. - -A mock adapter and a unit test that exercise this ABI on any Linux host live in -[`../../../tests/tpu/`](../../../tests/tpu/). + underlying PJRT buffer. The token is **not** the buffer's data and must never + be dereferenced — on real PJRT it is an internal handle that is + host-readable but reads as unrelated bytes. +- **Interior pointers:** `ProxyManager` stages a transfer in `chunk_size` + (4 MiB) pieces and passes `token + chunk_offset` for every chunk after the + first. Classification and copy entrypoints must therefore resolve an address + to the registered buffer whose range contains it. An adapter that only matches + base addresses makes TENT classify HBM as host memory, and — because the token + *is* readable — the staging copy degrades into a `memcpy` of unrelated bytes + that reports success. `TpuTransport` defends against this by requiring exactly + one TPU-device side per staging hop and failing loudly otherwise. + +A mock adapter and unit tests that exercise this ABI on any Linux host live in +[`../../../tests/tpu/`](../../../tests/tpu/). The mock hands out poisoned tokens +backed by shadow storage, so a copy that bypasses the adapter yields `0xDD` +rather than accidentally producing the right answer. ## Not yet included (follow-ups) diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt index 99b40e57e7..c7fd499ca8 100644 --- a/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt @@ -1,5 +1,8 @@ if(USE_TPU) file(GLOB XPORT_SOURCES "*.cpp") add_library(tent_xport_tpu STATIC ${XPORT_SOURCES}) - target_link_libraries(tent_xport_tpu PUBLIC tent_rpc tent_common) + # platform_tpu provides TpuPjrtShim, which the transport consults to + # verify that a staging copy really has a TPU-device side. + target_link_libraries(tent_xport_tpu PUBLIC tent_rpc tent_common + platform_tpu) endif() diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp index c3f7646497..380cf1e547 100644 --- a/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp @@ -17,6 +17,7 @@ #include #include "tent/common/status.h" +#include "tent/platform/tpu_pjrt_shim.h" #include "tent/runtime/platform.h" #include "tent/runtime/slab.h" @@ -59,7 +60,8 @@ Status TpuTransport::uninstall() { Status TpuTransport::allocateSubBatch(SubBatchRef &batch, size_t max_size) { auto tpu_batch = Slab::Get().allocate(); if (!tpu_batch) - return Status::InternalError("Unable to allocate TPU sub-batch"); + return Status::InternalError( + "Unable to allocate TPU sub-batch" LOC_MARK); batch = tpu_batch; tpu_batch->task_list.reserve(max_size); tpu_batch->max_size = max_size; @@ -108,6 +110,32 @@ void TpuTransport::startTransfer(TpuTask *task, TpuSubBatch *batch) { } void *staging = reinterpret_cast(task->request.target_offset); + + // Exactly one side of a staging hop is TPU HBM: the local stage copies + // HBM<->host staging buffer, and a delegated remote stage copies the peer's + // host staging buffer<->its HBM (so `staging` is the device side there). + // Verify that here instead of relying on Platform::copy to classify: if the + // adapter fails to recognise a device pointer -- e.g. it only matches base + // addresses and ProxyManager handed us `base + chunk_offset` -- then + // Platform::copy sees two host pointers and silently memcpy()s from a token + // that is not the buffer's data, corrupting the transfer. Fail loudly. + auto &shim = TpuPjrtShim::instance(); + const bool source_is_device = shim.isDevicePtr(task->request.source); + const bool staging_is_device = shim.isDevicePtr(staging); + if (source_is_device == staging_is_device) { + LOG(ERROR) + << "TpuTransport: a staging copy must have exactly one TPU " + "device side, but source=" + << task->request.source << " (device=" << source_is_device + << ") and target=" << staging << " (device=" << staging_is_device + << "). Either the PJRT adapter is unavailable, or it does not " + "resolve interior pointers (see tpu_pjrt_abi.h)."; + task->status_word = TransferStatusEnum::FAILED; + task->transferred_bytes = 0; + batch->notifyProgress(); + return; + } + Status status; if (task->request.opcode == Request::READ) // host staging buffer -> device (H2D) diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 7811769a42..c117b316b1 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -235,4 +235,18 @@ if(USE_TPU) tent_tpu_pjrt_shim_test PRIVATE MOCK_TPU_PJRT_LIB="$") add_test(NAME tent_tpu_pjrt_shim_test COMMAND tent_tpu_pjrt_shim_test) + + # Drives TpuTransport's staging hop (interior-offset chunks, and the + # exactly-one-device-side guard) against the same mock adapter. + add_executable(tent_tpu_transport_test tpu/tpu_transport_test.cpp) + add_dependencies(tent_tpu_transport_test mock_tpu_pjrt) + target_link_libraries(tent_tpu_transport_test + PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) + target_include_directories(tent_tpu_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_compile_definitions( + tent_tpu_transport_test + PRIVATE MOCK_TPU_PJRT_LIB="$") + add_test(NAME tent_tpu_transport_test COMMAND tent_tpu_transport_test) endif() diff --git a/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp index 48f569337b..0773faffd7 100644 --- a/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp +++ b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp @@ -12,24 +12,67 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Mock TPU/PJRT adapter for unit testing TpuPjrtShim and TpuPlatform without -// TPU hardware or a PJRT runtime. It implements the C ABI declared in -// tpu_pjrt_abi.h using ordinary host memory: "device" buffers are plain -// allocations the test registers via the mock_* test helpers below, and D2H/H2D -// copies are memcpy. This lets the routing and shim logic be exercised on any -// Linux host (see tpu_pjrt_shim_test.cpp). +// Mock TPU/PJRT adapter for unit testing TpuPjrtShim, TpuPlatform and +// TpuTransport without TPU hardware or a PJRT runtime. It implements the C ABI +// declared in tpu_pjrt_abi.h, so the routing, classification and staging logic +// can be exercised on any Linux host (see tpu_pjrt_shim_test.cpp and +// tpu_transport_test.cpp). +// +// Two properties of real PJRT/TPU are modelled deliberately, because both are +// load-bearing for TENT's correctness and neither is obvious: +// +// 1. The device "pointer" is an opaque TOKEN, not the buffer's data. On real +// hardware PJRT_Buffer_UnsafePointer returns an internal handle that is +// host-dereferenceable but reads as unrelated bytes. So the token here is a +// separate read-only mapping poisoned with kPoison, and the buffer contents +// live in a shadow allocation reachable only through the copy entrypoints. +// Any code path that "helpfully" memcpy()s from a device token therefore +// reads poison instead of silently producing the right answer -- which is +// what a plain-host-memory mock would have done, and why an +// interior-pointer bug could pass a green test suite. +// +// 2. Addresses may be INTERIOR to a registered buffer. TENT stages transfers +// in +// chunks and passes `token + chunk_offset` for every chunk after the first, +// so classification and copies resolve ranges, and a copy that would run +// past a buffer's end is rejected rather than truncated. + +#include +#include #include #include -#include +#include #include "tent/platform/tpu_pjrt_abi.h" namespace { +// Byte filling the token mapping. Nothing should ever read it; if a test sees +// 0xDD it means something dereferenced a device token directly. +constexpr unsigned char kPoison = 0xDD; + +struct Buffer { + uintptr_t token; // opaque handle handed out to TENT + unsigned char *shadow; // where the bytes actually live + size_t size; + int device; +}; + std::mutex g_mutex; -// Registered fake device buffers: pointer -> device ordinal. -std::unordered_map g_device_registry; +std::vector g_device_registry; int g_device_count = 4; + +// Returns the buffer containing [addr, addr + len), or nullptr. len == 0 only +// checks that `addr` itself is inside a buffer. +Buffer *findLocked(const void *addr, size_t len) { + auto a = reinterpret_cast(addr); + for (auto &b : g_device_registry) { + if (a < b.token || a >= b.token + b.size) continue; + if (len > b.token + b.size - a) return nullptr; // runs past the end + return &b; + } + return nullptr; +} } // namespace extern "C" { @@ -40,24 +83,32 @@ int mc_tpu_pjrt_init(void) { return 0; } int mc_tpu_pjrt_is_device_ptr(const void *addr) { std::lock_guard lock(g_mutex); - return g_device_registry.count(addr) ? 1 : 0; + return findLocked(addr, 0) ? 1 : 0; } int mc_tpu_pjrt_device_index(const void *addr) { std::lock_guard lock(g_mutex); - auto it = g_device_registry.find(addr); - return it == g_device_registry.end() ? -1 : it->second; + const Buffer *b = findLocked(addr, 0); + return b ? b->device : -1; } int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len) { if (!host_dst || !device_src) return 1; - std::memcpy(host_dst, device_src, len); + std::lock_guard lock(g_mutex); + const Buffer *b = findLocked(device_src, len); + if (!b) return 1; + size_t offset = reinterpret_cast(device_src) - b->token; + std::memcpy(host_dst, b->shadow + offset, len); return 0; } int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len) { if (!device_dst || !host_src) return 1; - std::memcpy(device_dst, host_src, len); + std::lock_guard lock(g_mutex); + const Buffer *b = findLocked(device_dst, len); + if (!b) return 1; + size_t offset = reinterpret_cast(device_dst) - b->token; + std::memcpy(b->shadow + offset, host_src, len); return 0; } @@ -71,13 +122,39 @@ int mc_tpu_pjrt_device_numa(int index) { // --- Test-only helpers (not part of the shim ABI) -------------------------- -void mock_tpu_pjrt_register_device(const void *addr, int index) { +// Creates a fake device buffer of `size` bytes on device `index` and returns +// its token. The token is readable (like PJRT's unsafe pointer) but holds +// poison, never the buffer's data. +void *mock_tpu_pjrt_register_device(size_t size, int index) { + void *token = mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (token == MAP_FAILED) return nullptr; + std::memset(token, kPoison, size); + mprotect(token, size, PROT_READ); + std::lock_guard lock(g_mutex); - g_device_registry[addr] = index; + g_device_registry.push_back(Buffer{reinterpret_cast(token), + new unsigned char[size](), size, index}); + return token; } +// Direct access to a buffer's bytes, for seeding inputs and asserting outputs. +// Accepts an interior token address and returns the matching shadow address. +void *mock_tpu_pjrt_device_data(const void *token) { + std::lock_guard lock(g_mutex); + Buffer *b = findLocked(token, 0); + if (!b) return nullptr; + return b->shadow + (reinterpret_cast(token) - b->token); +} + +unsigned char mock_tpu_pjrt_poison_byte(void) { return kPoison; } + void mock_tpu_pjrt_reset(void) { std::lock_guard lock(g_mutex); + for (auto &b : g_device_registry) { + munmap(reinterpret_cast(b.token), b.size); + delete[] b.shadow; + } g_device_registry.clear(); g_device_count = 4; } diff --git a/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp index e1fa7bc359..aa088274ca 100644 --- a/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp +++ b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp @@ -24,6 +24,7 @@ #include #include +#include #include #ifndef MOCK_TPU_PJRT_LIB @@ -34,7 +35,8 @@ namespace mooncake { namespace tent { namespace { -using RegisterFn = void (*)(const void *, int); +using RegisterFn = void *(*)(size_t, int); +using DeviceDataFn = void *(*)(const void *); using ResetFn = void (*)(); using SetCountFn = void (*)(int); @@ -53,10 +55,13 @@ class TpuShimTest : public ::testing::Test { ASSERT_NE(mock_, nullptr) << dlerror(); register_ = reinterpret_cast( dlsym(mock_, "mock_tpu_pjrt_register_device")); + device_data_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_device_data")); reset_ = reinterpret_cast(dlsym(mock_, "mock_tpu_pjrt_reset")); set_count_ = reinterpret_cast( dlsym(mock_, "mock_tpu_pjrt_set_device_count")); ASSERT_NE(register_, nullptr); + ASSERT_NE(device_data_, nullptr); ASSERT_NE(reset_, nullptr); ASSERT_NE(set_count_, nullptr); reset_(); @@ -67,8 +72,16 @@ class TpuShimTest : public ::testing::Test { if (mock_) dlclose(mock_); } + // Fills a fake device buffer with a repeatable pattern. + void seedDevice(void *token, size_t size, uint8_t modulus) { + auto *data = static_cast(device_data_(token)); + ASSERT_NE(data, nullptr); + for (size_t i = 0; i < size; ++i) data[i] = (uint8_t)(i % modulus); + } + void *mock_ = nullptr; RegisterFn register_ = nullptr; + DeviceDataFn device_data_ = nullptr; ResetFn reset_ = nullptr; SetCountFn set_count_ = nullptr; }; @@ -84,12 +97,12 @@ TEST_F(TpuShimTest, AdapterLoadsAndReportsDevices) { TEST_F(TpuShimTest, ClassifiesRegisteredDevicePointers) { int host_value = 0; - std::vector fake_device(64); - register_(fake_device.data(), /*index=*/2); + void *token = register_(64, /*index=*/2); + ASSERT_NE(token, nullptr); auto &shim = TpuPjrtShim::instance(); - EXPECT_TRUE(shim.isDevicePtr(fake_device.data())); - EXPECT_EQ(shim.deviceIndex(fake_device.data()), 2); + EXPECT_TRUE(shim.isDevicePtr(token)); + EXPECT_EQ(shim.deviceIndex(token), 2); // Unregistered host memory is not device memory. EXPECT_FALSE(shim.isDevicePtr(&host_value)); @@ -97,21 +110,79 @@ TEST_F(TpuShimTest, ClassifiesRegisteredDevicePointers) { EXPECT_FALSE(shim.isDevicePtr(nullptr)); } +// Regression: ProxyManager stages a transfer in chunk_size (4 MiB) pieces and +// hands the platform `token + chunk_offset` for every chunk after the first. An +// adapter that only recognises base addresses makes TENT classify TPU HBM as +// host memory, and the staging copy degrades into a memcpy from a token that on +// real PJRT is not the buffer's data -- silently corrupting every transfer +// larger than one chunk. Interior addresses must classify as device memory. +TEST_F(TpuShimTest, ClassifiesInteriorDevicePointers) { + const size_t kSize = 8192; + auto *token = static_cast(register_(kSize, /*index=*/3)); + ASSERT_NE(token, nullptr); + + auto &shim = TpuPjrtShim::instance(); + EXPECT_TRUE(shim.isDevicePtr(token + 1)); + EXPECT_TRUE(shim.isDevicePtr(token + 4096)); + EXPECT_TRUE(shim.isDevicePtr(token + kSize - 1)); + EXPECT_EQ(shim.deviceIndex(token + 4096), 3); + + // One past the end belongs to no buffer. + EXPECT_FALSE(shim.isDevicePtr(token + kSize)); + EXPECT_EQ(shim.deviceIndex(token + kSize), -1); +} + TEST_F(TpuShimTest, CopyRoundTripMovesBytes) { auto &shim = TpuPjrtShim::instance(); const std::vector src = {1, 2, 3, 4, 5, 6, 7, 8}; - std::vector device(src.size(), 0); + void *token = register_(src.size(), /*index=*/0); + ASSERT_NE(token, nullptr); std::vector dst(src.size(), 0); // host -> "device" - ASSERT_TRUE(shim.copyH2D(device.data(), src.data(), src.size()).ok()); - EXPECT_EQ(device, src); + ASSERT_TRUE(shim.copyH2D(token, src.data(), src.size()).ok()); + EXPECT_EQ(0, std::memcmp(device_data_(token), src.data(), src.size())); // "device" -> host - ASSERT_TRUE(shim.copyD2H(dst.data(), device.data(), device.size()).ok()); + ASSERT_TRUE(shim.copyD2H(dst.data(), token, src.size()).ok()); EXPECT_EQ(dst, src); } +// The chunked staging pattern end to end: copy each 4 KiB slice of a "device" +// buffer out through an interior pointer, exactly as ProxyManager would. +TEST_F(TpuShimTest, CopyFromInteriorOffsetMovesTheRightBytes) { + auto &shim = TpuPjrtShim::instance(); + const size_t kChunk = 4096, kChunks = 4, kSize = kChunk * kChunks; + auto *token = static_cast(register_(kSize, /*index=*/1)); + ASSERT_NE(token, nullptr); + seedDevice(token, kSize, 251); + + std::vector host(kSize, 0); + for (size_t c = 0; c < kChunks; ++c) { + ASSERT_TRUE( + shim.copyD2H(host.data() + c * kChunk, token + c * kChunk, kChunk) + .ok()) + << "chunk " << c; + } + EXPECT_EQ(0, std::memcmp(host.data(), device_data_(token), kSize)); + // And the bytes are the seeded pattern, not the token's poison. + EXPECT_EQ(host[0], 0); + EXPECT_EQ(host[kChunk + 1], (uint8_t)((kChunk + 1) % 251)); +} + +// A copy must never run past the end of the buffer it started in. +TEST_F(TpuShimTest, RejectsCopyRunningPastBufferEnd) { + auto &shim = TpuPjrtShim::instance(); + auto *token = static_cast(register_(1024, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector host(2048, 0); + + EXPECT_FALSE(shim.copyD2H(host.data(), token + 512, 1024).ok()); + EXPECT_FALSE(shim.copyH2D(token + 512, host.data(), 1024).ok()); + // Unregistered addresses are not device memory and cannot be copied. + EXPECT_FALSE(shim.copyD2H(host.data(), host.data() + 1024, 16).ok()); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp b/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp new file mode 100644 index 0000000000..ee92a45fc1 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp @@ -0,0 +1,204 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Drives TpuTransport's staging hop against the mock PJRT adapter. No TPU +// hardware or PJRT runtime required. +// +// The cases here mirror the requests ProxyManager actually issues: a local +// stage whose `source` is `device_token + chunk_offset` and whose +// `target_offset` is the host staging buffer, and the delegated remote stage +// where the device side is `target_offset` instead. +// +// The mock's device tokens are poisoned (see mock_tpu_pjrt_adapter.cpp), so a +// staging copy that bypasses the adapter and memcpy()s straight from a token +// yields 0xDD rather than the buffer's data. That makes the data assertions +// below real detectors of the "device pointer misclassified as host memory" +// failure mode, instead of accidentally passing. + +#include "tent/transport/tpu/tpu_transport.h" + +#include +#include + +#include +#include +#include + +#ifndef MOCK_TPU_PJRT_LIB +#error "MOCK_TPU_PJRT_LIB must be defined by the build (path to mock adapter)" +#endif + +namespace mooncake { +namespace tent { +namespace { + +using RegisterFn = void *(*)(size_t, int); +using DeviceDataFn = void *(*)(const void *); +using PoisonFn = unsigned char (*)(); +using ResetFn = void (*)(); + +constexpr size_t kChunk = 4ul << 20; // ProxyManager's default chunk_size +constexpr size_t kBufSize = 3 * kChunk; + +class TpuTransportTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + ::setenv("MC_TPU_PJRT_LIB", MOCK_TPU_PJRT_LIB, /*overwrite=*/1); + } + + void SetUp() override { + mock_ = dlopen(MOCK_TPU_PJRT_LIB, RTLD_NOW | RTLD_GLOBAL); + ASSERT_NE(mock_, nullptr) << dlerror(); + register_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_register_device")); + device_data_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_device_data")); + poison_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_poison_byte")); + reset_ = reinterpret_cast(dlsym(mock_, "mock_tpu_pjrt_reset")); + ASSERT_NE(register_, nullptr); + ASSERT_NE(device_data_, nullptr); + ASSERT_NE(poison_, nullptr); + ASSERT_NE(reset_, nullptr); + reset_(); + + std::string segment = "local"; + ASSERT_TRUE( + transport_.install(segment, nullptr, nullptr, nullptr).ok()); + } + + void TearDown() override { + transport_.uninstall(); + if (reset_) reset_(); + if (mock_) dlclose(mock_); + } + + // Runs one request through the transport and returns its final status. + TransferStatus run(const Request &request) { + Transport::SubBatchRef batch = nullptr; + EXPECT_TRUE(transport_.allocateSubBatch(batch, 1).ok()); + EXPECT_TRUE(transport_.submitTransferTasks(batch, {request}).ok()); + TransferStatus status{}; + EXPECT_TRUE(transport_.getTransferStatus(batch, 0, status).ok()); + EXPECT_TRUE(transport_.freeSubBatch(batch).ok()); + return status; + } + + static Request makeRequest(Request::OpCode op, void *source, + uint64_t target, size_t length) { + Request r; + r.opcode = op; + r.source = source; + r.length = length; + r.target_id = LOCAL_SEGMENT_ID; + r.target_offset = target; + return r; + } + + void *mock_ = nullptr; + RegisterFn register_ = nullptr; + DeviceDataFn device_data_ = nullptr; + PoisonFn poison_ = nullptr; + ResetFn reset_ = nullptr; + TpuTransport transport_; +}; + +// The local WRITE stage of chunk #1: source is an interior device address. +// Before interior pointers were part of the adapter contract this silently +// memcpy()d from a non-data token and still reported COMPLETED. +TEST_F(TpuTransportTest, LocalStageCopiesFromInteriorDeviceOffset) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + auto *data = static_cast(device_data_(token)); + for (size_t i = 0; i < kBufSize; ++i) data[i] = (uint8_t)(i % 251); + std::vector staging(kChunk, 0); + + // Chunk #1: device_token + 4 MiB -> host staging buffer. + auto status = run(makeRequest(Request::WRITE, token + kChunk, + (uint64_t)staging.data(), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(status.transferred_bytes, kChunk); + EXPECT_EQ(0, std::memcmp(staging.data(), data + kChunk, kChunk)); + // Not the token's poison: the adapter really did the copy. + EXPECT_NE(staging[0], poison_()); +} + +// The mirrored remote stage: the device side is `target_offset`, at an offset. +TEST_F(TpuTransportTest, RemoteStageCopiesToInteriorDeviceOffset) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk); + for (size_t i = 0; i < kChunk; ++i) staging[i] = (uint8_t)(i % 197); + + // WRITE with a device target: host staging -> device_token + 4 MiB. + auto status = run(makeRequest(Request::WRITE, staging.data(), + (uint64_t)(token + kChunk), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + auto *data = static_cast(device_data_(token)); + EXPECT_EQ(0, std::memcmp(data + kChunk, staging.data(), kChunk)); +} + +// A READ stage moves host staging -> device. +TEST_F(TpuTransportTest, ReadStageCopiesHostToDevice) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk, 0xAB); + + auto status = run(makeRequest(Request::READ, token + 2 * kChunk, + (uint64_t)staging.data(), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + auto *data = static_cast(device_data_(token)); + EXPECT_EQ(0, std::memcmp(data + 2 * kChunk, staging.data(), kChunk)); +} + +// Neither side is device memory: the adapter is missing, or it failed to +// classify an interior pointer. Either way the transport must fail rather than +// let Platform::copy memcpy from a token that is not the buffer's data. +TEST_F(TpuTransportTest, FailsWhenNeitherSideIsDeviceMemory) { + std::vector host_a(4096, 1), host_b(4096, 2); + auto status = run(makeRequest(Request::WRITE, host_a.data(), + (uint64_t)host_b.data(), host_a.size())); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); + EXPECT_EQ(status.transferred_bytes, 0u); + // The destination is untouched -- no silent partial copy. + EXPECT_EQ(host_b[0], 2); +} + +// Both sides device: HBM<->HBM is not a staging hop and must be rejected. +TEST_F(TpuTransportTest, FailsWhenBothSidesAreDeviceMemory) { + auto *a = static_cast(register_(kBufSize, /*index=*/0)); + auto *b = static_cast(register_(kBufSize, /*index=*/1)); + ASSERT_NE(a, nullptr); + ASSERT_NE(b, nullptr); + + auto status = run(makeRequest(Request::WRITE, a, (uint64_t)b, kChunk)); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); +} + +// TPU HBM can never be a remote peer; a non-local target is a routing bug. +TEST_F(TpuTransportTest, FailsOnNonLocalTarget) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk, 0); + + auto request = + makeRequest(Request::WRITE, token, (uint64_t)staging.data(), kChunk); + request.target_id = LOCAL_SEGMENT_ID + 1; + auto status = run(request); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); +} + +} // namespace +} // namespace tent +} // namespace mooncake From 3d0748cf7dc90cd385769c667fed080abbfdb9ab Mon Sep 17 00:00:00 2001 From: ykwd Date: Fri, 10 Jul 2026 10:24:39 +0800 Subject: [PATCH 43/57] [Doc] Reorganize performance docs (#2824) * clear legacy benchmark results * Change vllm performance benchmark docs menu * rename vllm performance file name * Change Mooncake performance docs paths * Change SGLang performance docs paths * update the sglang and vllm performance description --------- Co-authored-by: Ke Yang --- docs/source/deployment/ssd-offload.md | 2 +- docs/source/design/hicache-design.md | 4 +- docs/source/design/mooncake-store.md | 4 +- docs/source/design/ssd-offload.md | 2 +- .../examples/sglang-integration-v1.md | 2 +- .../hicache-integration-v1.md | 2 +- .../sglang-integration/hicache-quick-start.md | 2 +- .../examples/sglang-integration/index.md | 2 +- .../vllm-integration/disagg-prefill-decode.md | 4 +- .../examples/vllm-integration/index.md | 2 +- .../vllm-integration/kv-cache-storage.md | 11 +---- .../vllm-integration/vllm-integration-v0.2.md | 2 +- .../vllm-integration/vllm-integration-v1.0.md | 4 +- .../vllm-mooncakestoreconnector.md | 2 +- docs/source/index.md | 5 +- .../allocation-strategy-benchmark-result.md | 2 - .../allocator-benchmark-result.md | 0 docs/source/performance/mooncake/index.md | 18 +++---- .../ssd-offload-benchmark-results.md | 4 +- .../{ => mooncake}/storage-benchmark.md | 0 docs/source/performance/sglang/index.md | 12 ++--- .../sglang-benchmark-results-v1.md | 4 +- .../sglang-hicache-benchmark-results-v1.md | 10 ++-- .../vllm-benchmark-results-v0.2.md | 49 ------------------- .../performance/vllm-benchmark-results-v1.md | 22 --------- docs/source/performance/vllm/index.md | 19 +++---- .../vllm-v1-mooncake-store.md} | 3 +- .../vllm-v1-pd-performance.md} | 10 ++-- 28 files changed, 58 insertions(+), 145 deletions(-) rename docs/source/performance/{ => mooncake}/allocation-strategy-benchmark-result.md (99%) rename docs/source/performance/{ => mooncake}/allocator-benchmark-result.md (100%) rename docs/source/performance/{ => mooncake}/ssd-offload-benchmark-results.md (98%) rename docs/source/performance/{ => mooncake}/storage-benchmark.md (100%) rename docs/source/performance/{ => sglang}/sglang-benchmark-results-v1.md (96%) rename docs/source/performance/{ => sglang}/sglang-hicache-benchmark-results-v1.md (96%) delete mode 100644 docs/source/performance/vllm-benchmark-results-v0.2.md delete mode 100644 docs/source/performance/vllm-benchmark-results-v1.md rename docs/source/performance/{vllm-v1-kvcache-sharing.md => vllm/vllm-v1-mooncake-store.md} (98%) rename docs/source/performance/{vllm-v1-support-benchmark.md => vllm/vllm-v1-pd-performance.md} (90%) diff --git a/docs/source/deployment/ssd-offload.md b/docs/source/deployment/ssd-offload.md index d1ffe86b29..edf58f7008 100644 --- a/docs/source/deployment/ssd-offload.md +++ b/docs/source/deployment/ssd-offload.md @@ -4,7 +4,7 @@ Mooncake Store supports offloading KV cache objects from distributed memory to a local filesystem path, typically backed by local SSDs. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from the local filesystem-backed offload path. -For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md). +For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../performance/mooncake/ssd-offload-benchmark-results.md). SSD offload requires the **Real Client** and supports two deployment modes: diff --git a/docs/source/design/hicache-design.md b/docs/source/design/hicache-design.md index db2dfa2b2e..a2a301e869 100644 --- a/docs/source/design/hicache-design.md +++ b/docs/source/design/hicache-design.md @@ -2,7 +2,7 @@ With the rapid development of tasks such as Agentic Coding, the length of request contexts continues to grow. Increasing the capacity of the KV Cache to improve its hit rate has become increasingly important for enhancing throughput and reducing TTFT. In this context, SGLang introduces **HiCache**, which extends the original RadixAttention (previously limited to GPU memory) by adding hierarchical caching support and integrating with distributed storage backends such as Mooncake. -Inspired by the classic three-level cache design of modern CPUs, HiCache organizes GPU memory as L1, host memory as L2, and distributed storage as L3. This hierarchy enables HiCache to fully exploit the "idle" storage space of GPUs and CPUs, while integrating distributed cache systems for global KV cache storage and scheduling. As a result, HiCache significantly expands KV cache capacity while maintaining strong read performance, especially in workloads such as multi-QA and long-context inference, where KV cache reuse is frequent. For detailed benchmark results, see [this document](https://kvcache-ai.github.io/Mooncake/performance/sglang-hicache-benchmark-results-v1.html). +Inspired by the classic three-level cache design of modern CPUs, HiCache organizes GPU memory as L1, host memory as L2, and distributed storage as L3. This hierarchy enables HiCache to fully exploit the "idle" storage space of GPUs and CPUs, while integrating distributed cache systems for global KV cache storage and scheduling. As a result, HiCache significantly expands KV cache capacity while maintaining strong read performance, especially in workloads such as multi-QA and long-context inference, where KV cache reuse is frequent. For detailed benchmark results, see [this document](https://kvcache-ai.github.io/Mooncake/performance/sglang/sglang-hicache-benchmark-results-v1.html). While HiCache supports multiple L3 backends, this document focuses primarily on the **Mooncake** backend. @@ -115,7 +115,7 @@ Furthermore, **Mooncake** supports efficient batch read and write operations and ## Integration with PD-Disaggregation Deployment Mode -SGLang supports a PD (Prefill-Decode) disaggregation deployment mode through the **Mooncake TransferEngine** (for details, see [this document](https://docs.sglang.ai/advanced_features/pd_disaggregation.html)). +SGLang supports a PD (Prefill-Decode) disaggregation deployment mode through the **Mooncake TransferEngine** (for details, see [this document](https://docs.sglang.ai/advanced_features/pd_disaggregation.html)). In the PD-disaggregation deployment mode, HiCache can be enabled on the Prefill nodes to optimize prefill performance. With the hierarchical caching mechanism provided by **HiCache + Mooncake Store**, prefill nodes can handle long-context and multi-turn dialogue scenarios more efficiently, significantly improving performance during the prefill phase. HiCache can also be enabled on the decode nodes to write computation results back to L3. diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md index 6e02e3c8b9..bb34a1ad34 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -494,7 +494,7 @@ Mooncake Store provides two concrete implementations of `BufferAllocatorBase`: **OffsetBufferAllocator (default and recommended)**: This allocator is derived from [OffsetAllocator](https://github.com/sebbbi/OffsetAllocator), which uses a custom bin-based allocation strategy that supports fast hard realtime `O(1)` offset allocation with minimal fragmentation. Mooncake Store optimizes this allocator based on the specific memory usage characteristics of LLM inference workloads, thereby enhancing memory utilization in LLM scenarios. -For measured utilization and allocation latency across LLM-style workloads, see [Allocator Performance](../performance/allocator-benchmark-result.md). +For measured utilization and allocation latency across LLM-style workloads, see [Allocator Performance](../performance/mooncake/allocator-benchmark-result.md). **CachelibBufferAllocator (deprecated)**: This allocator leverages Facebook's [CacheLib](https://github.com/facebook/CacheLib) to manage memory using a slab-based allocation strategy. It provides efficient memory allocation with good fragmentation resistance and is well-suited for high-performance scenarios. However, in our modified version, it does not handle workloads with highly variable object sizes effectively, so it is currently marked as deprecated. @@ -588,7 +588,7 @@ Valid values are: `random` (default), `free_ratio_first`, `ssd_free_ratio_first` **Use `local_first`** when inference workers and Mooncake Store memory segments are colocated and you want writes to prefer the writer's host before falling back to other hosts. For this strategy to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique host part in `local_hostname`. -For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../performance/allocation-strategy-benchmark-result.md). +For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../performance/mooncake/allocation-strategy-benchmark-result.md). #### Strategy Details diff --git a/docs/source/design/ssd-offload.md b/docs/source/design/ssd-offload.md index d9430f5dd1..c55d1cd7e0 100644 --- a/docs/source/design/ssd-offload.md +++ b/docs/source/design/ssd-offload.md @@ -6,7 +6,7 @@ Mooncake Store supports offloading KV cache objects from distributed memory to l SSD offload is implemented as a background subsystem within the **real client** process. It is transparent to the application: a `Put` that would otherwise be evicted from memory is persisted to disk, and a `Get` that finds no memory replica automatically falls back to reading from SSD. -For multi-turn conversation benchmark results, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md). +For multi-turn conversation benchmark results, see [Mooncake SSD Offload Benchmark](../performance/mooncake/ssd-offload-benchmark-results.md). --- diff --git a/docs/source/getting_started/examples/sglang-integration-v1.md b/docs/source/getting_started/examples/sglang-integration-v1.md index 4338001dda..a169558bec 100644 --- a/docs/source/getting_started/examples/sglang-integration-v1.md +++ b/docs/source/getting_started/examples/sglang-integration-v1.md @@ -4,7 +4,7 @@ SGLang uses Mooncake's Transfer Engine to enable disaggregated prefill-decode (PD) serving across nodes via RDMA, with support for EP and EPD backends. This integration is based on [PR 4654](https://github.com/sgl-project/sglang/pull/4654) and [PR 4880](https://github.com/sgl-project/sglang/pull/4880). -In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput ([details](../../performance/sglang-benchmark-results-v1)). +In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput ([details](../../performance/sglang/sglang-benchmark-results-v1)). ``` +-----------+ Transfer Engine (RDMA) +-----------+ diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md b/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md index 4a6e011cb0..6e9383062b 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md +++ b/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md @@ -97,7 +97,7 @@ mooncake_master --enable_http_metadata_server=true --http_metadata_server_port=8 When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark ratio, an eviction task is triggered to free up space by evicting a portion of objects. -Due to memory fragmentation, allocation failures may occur even when memory usage has not yet reached 100%. The actual threshold depends on the workload. This [benchmark document](https://kvcache-ai.github.io/Mooncake/performance/allocator-benchmark-result.html) provides memory allocation efficiency results under different scenarios. if excessive allocation failures are observed, consider lowering this parameter accordingly. +Due to memory fragmentation, allocation failures may occur even when memory usage has not yet reached 100%. The actual threshold depends on the workload. This [benchmark document](https://kvcache-ai.github.io/Mooncake/performance/mooncake/allocator-benchmark-result.html) provides memory allocation efficiency results under different scenarios. if excessive allocation failures are observed, consider lowering this parameter accordingly. **Launch Mooncake `store service` (Optional):** diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md b/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md index a338ce2903..a4004f280f 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md +++ b/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md @@ -1,6 +1,6 @@ # Quick Start: SGLang HiCache with Mooncake Backend -Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend. In benchmarks, pre-populated Mooncake achieves **best TTFT** across all tiers, maintaining high cache hit rates as conversation rounds grow ([details](../../../performance/sglang-hicache-benchmark-results-v1)). +Follow this streamlined workflow to get SGLang HiCache running with Mooncake as the L3 storage backend. In benchmarks, pre-populated Mooncake achieves **best TTFT** across all tiers, maintaining high cache hit rates as conversation rounds grow ([details](../../../performance/sglang/sglang-hicache-benchmark-results-v1)). > Need more background or tuning options? See the [Complete Guide](hicache-integration-v1.md). diff --git a/docs/source/getting_started/examples/sglang-integration/index.md b/docs/source/getting_started/examples/sglang-integration/index.md index 885c444f2d..b927695bc7 100644 --- a/docs/source/getting_started/examples/sglang-integration/index.md +++ b/docs/source/getting_started/examples/sglang-integration/index.md @@ -17,7 +17,7 @@ SGLang uses Mooncake's Transfer Engine for direct zero-copy KV cache transfer be **Related:** [Full PD Disaggregation Guide](../sglang-integration-v1) — installation, cross-node/same-node setup, XpYd topology, EP backend for MoE models, and EPD backend for multimodal models. -**Benchmark:** [PD Disaggregation Performance](../../../performance/sglang-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances. +**Benchmark:** [PD Disaggregation Performance](../../../performance/sglang/sglang-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances. --- diff --git a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md index c8f0ab246d..49b81849f8 100644 --- a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md +++ b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md @@ -135,7 +135,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \ ### Performance -For detailed performance benchmarks and results, see the [vLLM Benchmark](../../../performance/vllm-v1-support-benchmark.md) documentation. +For detailed performance benchmarks and results, see the [vLLM PD Disaggregation Performance](../../../performance/vllm/vllm-v1-pd-performance.md) documentation. --- @@ -146,7 +146,7 @@ For detailed performance benchmarks and results, see the [vLLM Benchmark](../../ This section is for vLLM V0 backend (≤ v0.6.4.post1). For new deployments, use the [V1 backend](#using-vllm-v1-recommended) above. ``` -This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884). Preview benchmark results are available at [vLLM Benchmark Results V0.2](../../../performance/vllm-benchmark-results-v0.2.md). +This integration is based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884). ### Installation diff --git a/docs/source/getting_started/examples/vllm-integration/index.md b/docs/source/getting_started/examples/vllm-integration/index.md index f09f4c65b6..3c6d443c4f 100644 --- a/docs/source/getting_started/examples/vllm-integration/index.md +++ b/docs/source/getting_started/examples/vllm-integration/index.md @@ -5,7 +5,7 @@ Mooncake integrates with vLLM to accelerate large language model serving through high-performance KV cache transfer and shared storage. The integration supports two primary scenarios: - **Disaggregated Prefill-Decode Serving**: Seamlessly split prefill and decode across nodes using `MooncakeConnector`, with RDMA-powered cross-node KV cache transfer achieving up to **142.25 GB/s** peak bandwidth (71.1% utilization of 8x RoCE). Transfer overhead is negligible — for 32K-token prompts (4.50 GB of KV data), transfer takes only **31.65 ms**, accounting for just **4.2%** of total TTFT. -- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime. Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x** (1P1D, 12GPUs), and scales to **60 GPUs** with >95% cache hit rate as shown in this [webpage](../../../../../docs/source/performance/vllm-v1-kvcache-sharing.md). +- **KV Cache Storage & Sharing**: Extend effective KV cache capacity via `MooncakeStore` / `MooncakeStoreConnector`, with hash-based prefix caching that enables multiple vLLM instances to share cached KV blocks. Supports CPU/Disk offloading and dynamic XpYd topologies at runtime. Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x** (1P1D, 12GPUs), and scales to **60 GPUs** with >95% cache hit rate as shown in this [webpage](../../../performance/vllm/vllm-v1-mooncake-store.md). | Scenario | Guide | vLLM Backend | |----------|-------|-------------| diff --git a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md b/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md index c25427adef..c4d9128a6b 100644 --- a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md +++ b/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md @@ -4,7 +4,7 @@ This guide demonstrates how to use `MooncakeStore` / `MooncakeStoreConnector` with vLLM to build a distributed KV cache storage pool. It enables KV cache offloading to CPU/SSD, hash-based prefix caching across multiple vLLM instances, and flexible XpYd disaggregated deployment — where you can dynamically adjust prefill and decode group sizes at runtime. -Compared to Redis-based backends, MooncakeStore achieves significantly lower TTFT (e.g., **~32% improvement** in mean TTFT for 2P2D tp=2 under RDMA). See [benchmark results](../../../performance/vllm-benchmark-results-v1.md) for details. +Compared to Redis-based backends, MooncakeStore achieves significantly lower TTFT (e.g., **~32% improvement** in mean TTFT for 2P2D tp=2 under RDMA). --- @@ -388,15 +388,6 @@ curl -s http://localhost:8000/v1/completions \ --- -## Performance - -| Scenario | Document | -|----------|----------| -| V1 MooncakeStoreConnector vs Redis | [Benchmark V1](../../../performance/vllm-benchmark-results-v1.md) | -| V0 MooncakeStore vs Redis | [Benchmark V0](../../../performance/vllm-benchmark-results-v0.2.md) | - ---- - ## Troubleshooting - If you encounter connection issues, check that: diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md index c9421957fe..e4f41ce372 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md @@ -10,7 +10,7 @@ This page has been **consolidated** into the unified [Disaggregated Prefill-Deco ``` ## Overview -This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. We have run some experiments to obtain some [preview benchmark results](../../../performance/vllm-benchmark-results-v0.2.md). More benchmark results will be released in due time. +This is the latest version of mooncake-transfer-engine integration doc with the vLLM project based on [PR 10502](https://github.com/vllm-project/vllm/pull/10502) and [PR 10884](https://github.com/vllm-project/vllm/pull/10884) (vllm version: v0.6.4.post1/main) to accelerate KVCache transfer for inter-node disaggregated serving scenario. **_Please note that this is still an experimental version and will be modified anytime based on feedback from the vLLM community._** - **Update(Apr 10, 2025)**: We are working on the vLLM v1 integration now. Stay tuned. diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md index c1b2c31af7..86a6aedd5f 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md @@ -54,7 +54,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \ #### Proxy Server ```bash -# In vllm root directory. +# In vllm root directory. python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \ --prefiller-host 192.168.0.2 --prefiller-port 8010 \ --decoder-host 192.168.0.3 --decoder-port 8020 @@ -127,7 +127,7 @@ The following environment variables can be used to customize Mooncake behavior: ## Performance -For detailed performance benchmarks and results, see the [vLLM Benchmark](../../../performance/vllm-v1-support-benchmark.md) documentation. +For detailed performance benchmarks and results, see the [vLLM PD Disaggregation Performance](../../../performance/vllm/vllm-v1-pd-performance.md) documentation. ## Notes diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md index 25dff8fc4b..8517cb5de8 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md @@ -141,4 +141,4 @@ python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_ ### 4. Performance -Please refer to this [webpage](../../../../../docs/source/performance/vllm-v1-kvcache-sharing.md). \ No newline at end of file +Please refer to this [webpage](../../../performance/vllm/vllm-v1-mooncake-store.md). diff --git a/docs/source/index.md b/docs/source/index.md index 288b864f1c..5c7f6da659 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -111,10 +111,9 @@ Mooncake x LMDeploy Integration95% cache hit rate | -| [vLLM V1 + MooncakeConnector](../vllm-v1-support-benchmark) | vLLM V1 | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens | -| [vLLM V1 + MooncakeStore vs Redis](../vllm-benchmark-results-v1) | vLLM V1 | MooncakeStore RDMA consistently outperforms Redis across all XpYd topologies — e.g., **~32% lower** mean TTFT in 2P2D tp=2 | -| [vLLM V0 + MooncakeConnector (Legacy)](../vllm-benchmark-results-v0.2) | vLLM V0 | TP=4 reduces TTFT by ~80% vs TP=1; RDMA provides significant latency advantage over TCP across varying QPS and input lengths | +| Document | Scenario | Highlights | +|----------|----------|---------------| +| [PD Disaggregation Performance](vllm-v1-pd-performance) | PD disaggregation with Mooncake Connector | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens | +| [vLLM x Mooncake Store Performance](vllm-v1-mooncake-store) | distributed KV cache pool with Mooncake Store | Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x**, and scales to **60 GPUs** with >95% cache hit rate | :::{toctree} :maxdepth: 1 :hidden: -../vllm-v1-kvcache-sharing -../vllm-v1-support-benchmark -../vllm-benchmark-results-v1 -../vllm-benchmark-results-v0.2 +vllm-v1-pd-performance +vllm-v1-mooncake-store + ::: diff --git a/docs/source/performance/vllm-v1-kvcache-sharing.md b/docs/source/performance/vllm/vllm-v1-mooncake-store.md similarity index 98% rename from docs/source/performance/vllm-v1-kvcache-sharing.md rename to docs/source/performance/vllm/vllm-v1-mooncake-store.md index 8b63999ea9..df8a78c3db 100644 --- a/docs/source/performance/vllm-v1-kvcache-sharing.md +++ b/docs/source/performance/vllm/vllm-v1-mooncake-store.md @@ -1,4 +1,4 @@ -# Benchmark performance +# vLLM x Mooncake Store Performance Mooncake leverages the `MooncakeStoreConnector` in vLLM V1 to enable a distributed KV cache pool, supporting cross-instance sharing and reuse of KV caches. Furthermore, vLLM's `MultiConnector` can be configured to orchestrate both the `MooncakeConnector` (for peer-to-peer KV transfer) and the `MooncakeStoreConnector` (for the shared pool), enabling prefill-decode (PD) disaggregation. ![Overall Performance](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/hero_vllm_mooncake.svg) @@ -42,4 +42,3 @@ This result shows that the distributed KV cache pool substantially improves cach ## Benchmark Scripts The benchmark scripts are provided in the artifact repository [here](https://github.com/ivanium/vllm/tree/feat/mooncake-store-int/scripts/mooncake/artifacts). - diff --git a/docs/source/performance/vllm-v1-support-benchmark.md b/docs/source/performance/vllm/vllm-v1-pd-performance.md similarity index 90% rename from docs/source/performance/vllm-v1-support-benchmark.md rename to docs/source/performance/vllm/vllm-v1-pd-performance.md index d6cc7c578f..8780b3a4e5 100644 --- a/docs/source/performance/vllm-v1-support-benchmark.md +++ b/docs/source/performance/vllm/vllm-v1-pd-performance.md @@ -1,6 +1,6 @@ -# vLLM with Mooncake Transfer Engine Benchmark +# vLLM PD Disaggregation Performance -Mooncake has now implemented a vLLM connector, enabling direct support for the Prefill-Decode (PD) separation architecture in vLLM v1. We evaluated the performance of this integration, focusing on the efficiency of cross-node KV cache transfer using RDMA. +Mooncake has now implemented a vLLM connector, enabling direct support for the Prefill-Decode (PD) disaggregation architecture in vLLM v1. We evaluated the performance of this integration, focusing on the efficiency of cross-node KV cache transfer using RDMA. ## Benchmark Result @@ -8,7 +8,7 @@ Mooncake has now implemented a vLLM connector, enabling direct support for the P We measured the actual transfer bandwidth during the execution of requests with varying prompt lengths. -![KV Transfer Bandwidth (Actual)](../image/vllm_benchmark_actual_bandwidth.png) +![KV Transfer Bandwidth (Actual)](../../image/vllm_benchmark_actual_bandwidth.png) In a 1P1D (1 Prefiller, 1 Decoder) configuration using the Qwen3-8B model, Mooncake achieved a peak actual transfer bandwidth of **142.25 GB/s**. Given the theoretical maximum bandwidth of approximately 200 GB/s for the 8x RoCE connections, this represents a **71.1% bandwidth utilization rate**. This efficiency demonstrates that the custom transfer protocol and GPU Direct RDMA capabilities can effectively saturate high-performance networks. @@ -16,9 +16,9 @@ In a 1P1D (1 Prefiller, 1 Decoder) configuration using the Qwen3-8B model, Moonc We analyzed the Time To First Token (TTFT) to understand the impact of KV transfer overhead on end-to-end latency. -![TTFT Breakdown](../image/vllm_benchmark_ttft_breakdown.png) +![TTFT Breakdown](../../image/vllm_benchmark_ttft_breakdown.png) -![Transfer Time vs KV Size](../image/vllm_benchmark_transfer_time.png) +![Transfer Time vs KV Size](../../image/vllm_benchmark_transfer_time.png) The results show that Mooncake's high-speed transfer ensures that the overhead of moving KV cache is negligible compared to the computation time. For a prompt length of 32,768 tokens (transferring 4.50 GB of data), the actual KV transfer took only **31.65 ms**, accounting for merely **4.2%** of the total TTFT. From 0b4a12c5a77136ccb37b43f4cc4569fbfb12779d Mon Sep 17 00:00:00 2001 From: Csrayz <33659823+Csrayz@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:32:43 +0800 Subject: [PATCH 44/57] [Bugfix] Skip os.chmod when binary is already readable and executable (#2803) Signed-off-by: Csrayz <33659823+Csrayz@users.noreply.github.com> --- mooncake-wheel/mooncake/cli.py | 9 ++++++--- mooncake-wheel/mooncake/cli_bench.py | 9 ++++++--- mooncake-wheel/mooncake/cli_client.py | 5 ++++- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/mooncake-wheel/mooncake/cli.py b/mooncake-wheel/mooncake/cli.py index e252d5a313..2fba48a38f 100644 --- a/mooncake-wheel/mooncake/cli.py +++ b/mooncake-wheel/mooncake/cli.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -16,10 +17,12 @@ def main(): # Get the path to the mooncake_master binary package_dir = os.path.dirname(os.path.abspath(__file__)) bin_path = os.path.join(package_dir, "mooncake_master") - + # Make sure the binary is executable - os.chmod(bin_path, 0o755) - + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) diff --git a/mooncake-wheel/mooncake/cli_bench.py b/mooncake-wheel/mooncake/cli_bench.py index 2eb5ea986b..f4bf27d766 100644 --- a/mooncake-wheel/mooncake/cli_bench.py +++ b/mooncake-wheel/mooncake/cli_bench.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -16,10 +17,12 @@ def main(): # Get the path to the transfer_engine_bench binary package_dir = os.path.dirname(os.path.abspath(__file__)) bin_path = os.path.join(package_dir, "transfer_engine_bench") - + # Make sure the binary is executable - os.chmod(bin_path, 0o755) - + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) diff --git a/mooncake-wheel/mooncake/cli_client.py b/mooncake-wheel/mooncake/cli_client.py index c32cefebe9..0cf20ec717 100644 --- a/mooncake-wheel/mooncake/cli_client.py +++ b/mooncake-wheel/mooncake/cli_client.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -18,7 +19,9 @@ def main(): bin_path = os.path.join(package_dir, "mooncake_client") # Make sure the binary is executable - os.chmod(bin_path, 0o755) + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) From 8dee809219fa99ad429a4348a20478becf5b741c Mon Sep 17 00:00:00 2001 From: Aoi Date: Fri, 10 Jul 2026 11:27:08 +0800 Subject: [PATCH 45/57] [CI/Build] Fix compile warnings (#2825) --- mooncake-store/src/file_storage.cpp | 4 ++-- mooncake-store/src/master_metric_manager.cpp | 2 +- mooncake-store/src/master_service.cpp | 6 +++--- mooncake-store/src/real_client.cpp | 16 +++++++++------- mooncake-store/tests/mmap_arena_test.cpp | 5 +++-- .../rpc_communicator/rpc_communicator.h | 5 +++-- .../rpc_communicator/rpc_communicator.cpp | 15 ++++++++++----- .../transport/rpc_communicator/rpc_interface.cpp | 2 +- 8 files changed, 32 insertions(+), 23 deletions(-) diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 58bb601aaa..05d7ef2e02 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -446,8 +446,8 @@ tl::expected FileStorage::OffloadObjects( } std::unordered_map> user_batch_object; - auto query_result = BatchQuerySegmentSlices(user_keys, tenant_id, - user_batch_object); + [[maybe_unused]] auto query_result = BatchQuerySegmentSlices( + user_keys, tenant_id, user_batch_object); // BatchQuerySegmentSlices is now best-effort: it always returns // OK. Keys present in user_batch_object go to batch_object; the // rest are reported as failed. diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index d9c57cd4e0..ac4a38bfa4 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -1937,7 +1937,7 @@ std::string MasterMetricManager::get_summary_string( int64_t nof_allocated = nof_allocated_size_.value(); int64_t nof_capacity = nof_total_capacity_.value(); int64_t file_allocated = file_allocated_size_.value(); - int64_t file_capacity = file_total_capacity_.value(); + [[maybe_unused]] int64_t file_capacity = file_total_capacity_.value(); int64_t keys = key_count_.value(); int64_t soft_pin_keys = soft_pin_key_count_.value(); int64_t active_clients = active_clients_.value(); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index f3c2403dc5..edeeefe184 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -235,12 +235,12 @@ MasterService::MasterService(const MasterServiceConfig& config) config.snapshot_catalog_store_connstring), put_start_discard_timeout_sec_(config.put_start_discard_timeout_sec), put_start_release_timeout_sec_(config.put_start_release_timeout_sec), - task_manager_(config.task_manager_config), cxl_path_(config.cxl_path), cxl_size_(config.cxl_size), enable_cxl_(config.enable_cxl), offloading_queue_limit_(config.offloading_queue_limit), - offload_cap_ratio_(config.offload_cap_ratio) { + offload_cap_ratio_(config.offload_cap_ratio), + task_manager_(config.task_manager_config) { // Initialize HTTP metadata key prefix (read env var once at startup) const char* custom_prefix = std::getenv("MC_METADATA_CLUSTER_ID"); if (custom_prefix && std::strlen(custom_prefix) > 0) { @@ -4502,7 +4502,7 @@ auto MasterService::Remove(const std::string& key, const std::string& tenant_id, } PublishKvRemoved(key, metadata, tenant_id); - auto& tenant_state = accessor.GetTenantState(); + auto& tenant_state [[maybe_unused]] = accessor.GetTenantState(); accessor.Erase(); return {}; } diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 2f35506e5a..fe00d17224 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -4377,7 +4377,7 @@ std::vector> RealClient::batch_get_into_internal(const std::vector &keys, const std::vector &buffers, const std::vector &sizes) { - auto start_time = std::chrono::steady_clock::now(); + [[maybe_unused]] auto start_time = std::chrono::steady_clock::now(); // Validate preconditions if (!client_) { LOG(ERROR) << "Client is not initialized"; @@ -4646,8 +4646,9 @@ RealClient::batch_get_into_internal(const std::vector &keys, store_segment_it->second.emplace(op_it.first, op_it.second.slices); } - size_t offload_object_count = 0; - auto start_read_store_time = std::chrono::steady_clock::now(); + [[maybe_unused]] size_t offload_object_count = 0; + [[maybe_unused]] auto start_read_store_time = + std::chrono::steady_clock::now(); for (auto &offload_objects_it : offload_objects) { offload_object_count += offload_objects_it.second.size(); auto batch_get_offload_result = batch_get_into_offload_object_internal( @@ -4664,10 +4665,11 @@ RealClient::batch_get_into_internal(const std::vector &keys, } auto end_time = std::chrono::steady_clock::now(); - auto elapsed_time = std::chrono::duration_cast( - end_time - start_time) - .count(); - auto read_store_time = + [[maybe_unused]] auto elapsed_time = + std::chrono::duration_cast(end_time - + start_time) + .count(); + [[maybe_unused]] auto read_store_time = std::chrono::duration_cast( end_time - start_read_store_time) .count(); diff --git a/mooncake-store/tests/mmap_arena_test.cpp b/mooncake-store/tests/mmap_arena_test.cpp index 05425f0bd2..438aea2b14 100644 --- a/mooncake-store/tests/mmap_arena_test.cpp +++ b/mooncake-store/tests/mmap_arena_test.cpp @@ -651,10 +651,11 @@ TEST_F(MmapArenaTest, PagesArePhysicallyBackedAfterInit) { << "), falling back to read-verification"; // Read every page — if MAP_POPULATE didn't work, this would trigger // page faults (which is fine for CPU but would crash GPU DMA). - volatile char sink = 0; + char sum = 0; for (size_t off = 0; off < pool_size; off += sys_page_size) { - sink += static_cast(base)[off]; + sum += static_cast(base)[off]; } + volatile char sink = sum; (void)sink; // If we get here without SIGSEGV, at least CPU access works. // The real MAP_POPULATE guarantee is that DMA works too, which diff --git a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h index 88733b53e4..bafd3a09db 100644 --- a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h +++ b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h @@ -74,9 +74,10 @@ class RpcCommunicator { std::unique_ptr server_; std::function data_receive_callback_; - pybind11::handle py_callback_; + struct PyCallbackHolder; + std::unique_ptr py_callback_; std::shared_ptr> client_pools_; }; -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp index 158725f0ed..4cc7e93a38 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp @@ -15,7 +15,7 @@ namespace mooncake { namespace py = pybind11; -class py_rpc_context { +class __attribute__((visibility("hidden"))) py_rpc_context { public: void response_msg(py::buffer msg, py::object done) { py::buffer_info info = msg.request(); @@ -33,7 +33,12 @@ class py_rpc_context { coro_rpc::context context_; }; -RpcCommunicator::RpcCommunicator() {} +struct __attribute__((visibility("hidden"))) RpcCommunicator::PyCallbackHolder { + py::handle callback; +}; + +RpcCommunicator::RpcCommunicator() + : py_callback_(std::make_unique()) {} RpcCommunicator::~RpcCommunicator() { stopServer(); } @@ -395,7 +400,7 @@ void RpcCommunicator::handleDataTransferWithAttachment( auto view = py::memoryview::from_buffer(data.data(), {data.size()}, {sizeof(char)}); - py_callback_(std::move(t), view); + py_callback_->callback(std::move(t), view); } void RpcCommunicator::handleTensorTransferWithAttachment( @@ -412,7 +417,7 @@ void RpcCommunicator::handleTensorTransferWithAttachment( auto view = py::memoryview::from_buffer( attachment.data(), {attachment.size()}, {sizeof(int8_t)}); - py_callback_(std::move(t), view); + py_callback_->callback(std::move(t), view); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp index 13d44f6ff7..28b6dab51f 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp @@ -16,7 +16,7 @@ static constexpr size_t MAX_TENSOR_DIMS = 4; static constexpr size_t TENSOR_METADATA_SIZE = 4 + 4 + MAX_TENSOR_DIMS * 8; // Implementation class -class RpcInterface::Impl { +class __attribute__((visibility("hidden"))) RpcInterface::Impl { public: std::unique_ptr communicator; pybind11::function data_receive_callback; From 98ff4e4787e99265d25938139551841350ca5f4e Mon Sep 17 00:00:00 2001 From: xiangui <120565419+xiangui33423@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:35:57 +0800 Subject: [PATCH 46/57] [Store] Refactor: Extract snapshot orchestration into MasterSnapshotManager (#2805) Co-authored-by: Claude Opus 4.8 (1M context) --- mooncake-store/include/master_service.h | 46 +- .../include/master_snapshot_manager.h | 132 ++++ .../include/master_snapshot_repository.h | 92 +++ mooncake-store/src/CMakeLists.txt | 2 + mooncake-store/src/master_service.cpp | 743 +----------------- .../src/master_snapshot_manager.cpp | 712 +++++++++++++++++ .../src/master_snapshot_repository.cpp | 142 ++++ .../master_service_test_for_snapshot_base.h | 39 +- .../snapshot/snapshot_child_process_test.cpp | 97 ++- 9 files changed, 1250 insertions(+), 755 deletions(-) create mode 100644 mooncake-store/include/master_snapshot_manager.h create mode 100644 mooncake-store/include/master_snapshot_repository.h create mode 100644 mooncake-store/src/master_snapshot_manager.cpp create mode 100644 mooncake-store/src/master_snapshot_repository.cpp diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 661109fd51..42fd27efa3 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -39,6 +39,10 @@ #include "kv_event/kv_event_publisher.h" namespace mooncake { + +// Forward declaration for MasterSnapshotManager +class MasterSnapshotManager; + namespace ha { class SnapshotCatalogStore; } @@ -88,6 +92,8 @@ class MasterService { friend class test::PromotionOnHitTest; friend class benchmarks::BatchEvictBench; friend class test::MasterServiceTenantQuotaTest; + friend class MasterSnapshotManager; // Allow access to internal state for + // snapshot public: using NoFProbeFn = @@ -795,30 +801,7 @@ class MasterService { void setHttpMetadataRemoteUrl(const std::string& metadata_connstring); private: - void SnapshotThreadFunc(); - - // Persist master state - tl::expected PersistState( - const std::string& snapshot_id); - tl::expected PersistState( - const ha::SnapshotDescriptor& descriptor); - tl::expected - BuildSnapshotDescriptor(const std::string& snapshot_id, - const std::string& manifest_path, - const std::string& object_prefix) const; - tl::expected - ResolveSnapshotSequenceId() const; -#ifdef STORE_USE_ETCD - tl::expected - GetSnapshotBoundaryOpLogStore() const; -#endif - - tl::expected UploadSnapshotPayloadFile( - const std::vector& data, const std::string& path, - const std::string& local_filename, const std::string& snapshot_id); - std::unique_ptr CreateSnapshotCatalogStore(); - void CleanupOldSnapshot(int keep_count, const std::string& snapshot_id); ha::SnapshotCatalogStore* GetSnapshotCatalogStore(); // Restore master state @@ -828,12 +811,6 @@ class MasterService { const std::chrono::system_clock::time_point& now); void ResetStateAfterFailedRestoreAttempt(); - void WaitForSnapshotChild(pid_t pid, const std::string& snapshot_id, - int log_pipe_fd); - - void HandleChildTimeout(pid_t pid, const std::string& snapshot_id); - void HandleChildExit(pid_t pid, int status, const std::string& snapshot_id); - // BatchEvict evicts objects in a near-LRU way, i.e., prioritizes to evict // object with smaller lease timeout. It has two passes. The first pass only // evicts objects without soft pin. The second pass prioritizes objects @@ -1579,10 +1556,9 @@ class MasterService { static constexpr uint64_t kEvictionThreadSleepMs = 10; // 10 ms sleep between eviction checks - std::thread snapshot_thread_; - std::atomic snapshot_running_{false}; - std::mutex snapshot_thread_mutex_; - std::condition_variable snapshot_thread_cv_; + // Snapshot manager handles snapshot lifecycle orchestration + std::unique_ptr snapshot_manager_; + // Task cleanup thread related members std::thread task_cleanup_thread_; std::atomic task_cleanup_running_{false}; @@ -2026,10 +2002,6 @@ class MasterService { std::unique_ptr snapshot_object_store_; std::unique_ptr snapshot_catalog_store_; mutable std::shared_mutex snapshot_mutex_; -#ifdef STORE_USE_ETCD - mutable std::mutex snapshot_boundary_oplog_store_mutex_; - mutable std::unique_ptr snapshot_boundary_oplog_store_; -#endif // Discarded replicas management const std::chrono::seconds put_start_discard_timeout_sec_; diff --git a/mooncake-store/include/master_snapshot_manager.h b/mooncake-store/include/master_snapshot_manager.h new file mode 100644 index 0000000000..23e3f62626 --- /dev/null +++ b/mooncake-store/include/master_snapshot_manager.h @@ -0,0 +1,132 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "types.h" +#include "ha/ha_types.h" + +namespace mooncake { + +// Forward declarations +class MasterService; +class SnapshotObjectStore; +class MasterSnapshotRepository; + +namespace ha { +class SnapshotCatalogStore; +} + +namespace test { +class MasterServiceSnapshotTestBase; +class SnapshotChildProcessTest; +} // namespace test + +#ifdef STORE_USE_ETCD +class EtcdOpLogStore; +#endif + +struct MasterSnapshotManagerOptions { + bool enable_snapshot{false}; + uint64_t snapshot_interval_seconds{0}; + uint64_t snapshot_child_timeout_seconds{0}; + uint32_t snapshot_retention_count{0}; + std::string snapshot_backup_dir; + bool use_snapshot_backup_dir{false}; + std::string snapshot_catalog_store_type; + std::string snapshot_catalog_store_connstring; + std::string ha_backend_type; + std::string ha_backend_connstring; + std::string cluster_id; + bool enable_ha{false}; +}; + +/** + * @brief MasterSnapshotManager handles snapshot lifecycle orchestration for + * MasterService. This includes periodic snapshot scheduling, snapshot ID + * generation, descriptor construction, child process lifecycle management, + * timeout handling, payload upload, catalog publish, retention cleanup, and + * snapshot metrics updates. + * + * This is a behavior-preserving refactor that moves snapshot orchestration + * logic out of MasterService without changing snapshot format, restore + * behavior, storage layout, flags, or locking semantics. + */ +class MasterSnapshotManager { + friend class test::MasterServiceSnapshotTestBase; // Allow test access to + // private methods + friend class test::SnapshotChildProcessTest; // Allow test access to + // private methods + + public: + MasterSnapshotManager(MasterService* master_service, + MasterSnapshotManagerOptions options, + std::shared_mutex& snapshot_mutex, + SnapshotObjectStore* snapshot_object_store, + ha::SnapshotCatalogStore* snapshot_catalog_store); + + ~MasterSnapshotManager(); + + void Start(); + void Stop(); + + private: + void SnapshotThreadFunc(); + void WaitForSnapshotChild(pid_t pid, const std::string& snapshot_id, + int log_pipe_fd); + void HandleChildTimeout(pid_t pid, const std::string& snapshot_id); + void HandleChildExit(pid_t pid, int status, const std::string& snapshot_id); + + tl::expected PersistState( + const std::string& snapshot_id); + tl::expected PersistState( + const ha::SnapshotDescriptor& descriptor); + tl::expected + BuildSnapshotDescriptor(const std::string& snapshot_id, + const std::string& manifest_path, + const std::string& object_prefix) const; + tl::expected + ResolveSnapshotSequenceId() const; + +#ifdef STORE_USE_ETCD + tl::expected + GetSnapshotBoundaryOpLogStore() const; +#endif + + tl::expected UploadSnapshotPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id); + + void CleanupOldSnapshot(size_t keep_count, const std::string& snapshot_id); + + std::string FormatTimestamp( + const std::chrono::system_clock::time_point& tp); + + MasterService* master_service_; + MasterSnapshotManagerOptions options_; + + std::shared_mutex& snapshot_mutex_; + SnapshotObjectStore* snapshot_object_store_; + ha::SnapshotCatalogStore* snapshot_catalog_store_; + + std::unique_ptr repository_; + +#ifdef STORE_USE_ETCD + mutable std::mutex snapshot_boundary_oplog_store_mutex_; + mutable std::unique_ptr snapshot_boundary_oplog_store_; +#endif + + std::thread snapshot_thread_; + std::atomic snapshot_running_{false}; + std::mutex snapshot_thread_mutex_; + std::condition_variable snapshot_thread_cv_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/master_snapshot_repository.h b/mooncake-store/include/master_snapshot_repository.h new file mode 100644 index 0000000000..215e3c136b --- /dev/null +++ b/mooncake-store/include/master_snapshot_repository.h @@ -0,0 +1,92 @@ +#pragma once + +#include +#include + +#include + +#include "types.h" +#include "ha/ha_types.h" + +namespace mooncake { + +// Forward declarations +class SnapshotObjectStore; + +namespace ha { +class SnapshotCatalogStore; +} + +/** + * @brief MasterSnapshotRepository handles storage and catalog operations for + * snapshots. This includes uploading payload files to object storage, + * publishing snapshots to the catalog, listing snapshots, deleting snapshots, + * and enforcing retention policies. + * + * This class encapsulates all interactions with SnapshotObjectStore and + * SnapshotCatalogStore, separating storage concerns from snapshot orchestration + * logic in MasterSnapshotManager. + */ +class MasterSnapshotRepository { + public: + MasterSnapshotRepository(SnapshotObjectStore* object_store, + ha::SnapshotCatalogStore* catalog_store, + const std::string& backup_dir, + bool use_backup_dir); + + /** + * @brief Upload a single snapshot payload file to object storage + * @param data Binary data to upload + * @param path Storage path/key + * @param local_filename Filename for logging and local backup + * @param snapshot_id Snapshot ID for logging + * @return Empty on success, SerializationError on failure + */ + tl::expected UploadPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id); + + /** + * @brief Publish snapshot descriptor to catalog store + * @param descriptor Snapshot descriptor to publish + * @return ErrorCode::OK on success, error code on failure + */ + ErrorCode PublishSnapshot(const ha::SnapshotDescriptor& descriptor); + + /** + * @brief Cleanup old snapshots based on retention policy + * @param keep_count Number of recent snapshots to keep + * @param current_snapshot_id Current snapshot ID (for logging) + */ + void CleanupOldSnapshots(size_t keep_count, + const std::string& current_snapshot_id); + + /** + * @brief List all snapshots from catalog store + * @param limit Maximum number of snapshots to return (0 = unlimited) + * @return Vector of snapshot descriptors on success, error code on failure + */ + tl::expected, ErrorCode> ListSnapshots( + size_t limit); + + /** + * @brief Delete a specific snapshot from catalog store + * @param snapshot_id Snapshot ID to delete + * @return ErrorCode::OK on success, error code on failure + */ + ErrorCode DeleteSnapshot(const ha::SnapshotId& snapshot_id); + + /** + * @brief Get object store connection info for logging + * @return Connection info string + */ + std::string GetObjectStoreConnectionInfo() const; + + private: + SnapshotObjectStore* object_store_; + ha::SnapshotCatalogStore* catalog_store_; + std::string backup_dir_; + bool use_backup_dir_; +}; + +} // namespace mooncake diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index f85b42c280..63a4b502fc 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -4,6 +4,8 @@ add_subdirectory(cachelib_memory_allocator) set(MOONCAKE_STORE_SOURCES allocator.cpp master_service.cpp + master_snapshot_manager.cpp + master_snapshot_repository.cpp client_service.cpp client_metric.cpp types.cpp diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index edeeefe184..bf55656980 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -45,6 +45,7 @@ #include "utils/file_util.h" #include "utils.h" #include "kv_event/kv_event_config.h" +#include "master_snapshot_manager.h" namespace mooncake { @@ -458,9 +459,30 @@ MasterService::MasterService(const MasterServiceConfig& config) if (enable_snapshot_) { if (memory_allocator_type_ == BufferAllocatorType::OFFSET) { - snapshot_running_ = true; - snapshot_thread_ = - std::thread(&MasterService::SnapshotThreadFunc, this); + // Initialize and start snapshot manager + MasterSnapshotManagerOptions snapshot_options; + snapshot_options.enable_snapshot = enable_snapshot_; + snapshot_options.snapshot_interval_seconds = + snapshot_interval_seconds_; + snapshot_options.snapshot_child_timeout_seconds = + snapshot_child_timeout_seconds_; + snapshot_options.snapshot_retention_count = + snapshot_retention_count_; + snapshot_options.snapshot_backup_dir = snapshot_backup_dir_; + snapshot_options.use_snapshot_backup_dir = use_snapshot_backup_dir_; + snapshot_options.snapshot_catalog_store_type = + snapshot_catalog_store_type_; + snapshot_options.snapshot_catalog_store_connstring = + snapshot_catalog_store_connstring_; + snapshot_options.ha_backend_type = ha_backend_type_; + snapshot_options.ha_backend_connstring = ha_backend_connstring_; + snapshot_options.cluster_id = cluster_id_; + snapshot_options.enable_ha = enable_ha_; + + snapshot_manager_ = std::make_unique( + this, snapshot_options, snapshot_mutex_, + snapshot_object_store_.get(), snapshot_catalog_store_.get()); + snapshot_manager_->Start(); } } @@ -511,10 +533,12 @@ MasterService::~MasterService() { // Stop and join the threads eviction_running_ = false; client_monitor_running_ = false; - { - std::lock_guard lk(snapshot_thread_mutex_); - snapshot_running_ = false; + + // Stop snapshot manager (non-blocking) + if (snapshot_manager_) { + snapshot_manager_->Stop(); } + task_cleanup_running_ = false; job_dispatch_running_ = false; http_metadata_cleanup_running_ = false; @@ -524,7 +548,6 @@ MasterService::~MasterService() { #endif // Wake sleepers so join() doesn't block for long sleep intervals. - snapshot_thread_cv_.notify_all(); task_cleanup_cv_.notify_all(); http_metadata_cleanup_cv_.notify_all(); @@ -539,9 +562,6 @@ MasterService::~MasterService() { nof_heartbeat_thread_.join(); } #endif - if (snapshot_thread_.joinable()) { - snapshot_thread_.join(); - } if (task_cleanup_thread_.joinable()) { task_cleanup_thread_.join(); } @@ -551,6 +571,12 @@ MasterService::~MasterService() { if (job_dispatch_thread_.joinable()) { job_dispatch_thread_.join(); } + + // Reset snapshot manager after all other threads have joined + // This triggers the destructor which joins the snapshot thread + if (snapshot_manager_) { + snapshot_manager_.reset(); + } } void MasterService::SetNoFProbeFnForTesting(NoFProbeFn fn) { @@ -5875,686 +5901,6 @@ uint64_t MasterService::ReleaseExpiredDiscardedReplicas( return released_cnt; } -void MasterService::SnapshotThreadFunc() { - LOG(INFO) << "[Snapshot] snapshot_thread started"; - while (snapshot_running_) { - // Wait for the next snapshot cycle, but allow fast shutdown. - { - std::unique_lock lk(snapshot_thread_mutex_); - snapshot_thread_cv_.wait_for( - lk, std::chrono::seconds(snapshot_interval_seconds_), - [&] { return !snapshot_running_.load(); }); - } - - if (!snapshot_running_) { - break; - } - - if (!enable_snapshot_) { - // Snapshot is disabled - LOG(INFO) - << "[Snapshot] Snapshot is disabled, waiting for next cycle"; - continue; - } - // Fork a child process to save current state - - std::string snapshot_id = - FormatTimestamp(std::chrono::system_clock::now()); - LOG(INFO) << "[Snapshot] Preparing to fork child process, snapshot_id=" - << snapshot_id; - - // Create pipe for child process logging - int log_pipe[2]; - if (pipe(log_pipe) == -1) { - LOG(ERROR) << "[Snapshot] Failed to create log pipe: " - << strerror(errno) << ", snapshot_id=" << snapshot_id; - continue; - } - - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - const std::string path_prefix = snapshot_root + snapshot_id + "/"; - const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - auto descriptor = - BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); - if (!descriptor) { - LOG(ERROR) << "[Snapshot] Failed to build descriptor before fork, " - "snapshot_id=" - << snapshot_id - << ", code=" << toString(descriptor.error().code) - << ", msg=" << descriptor.error().message; - close(log_pipe[0]); - close(log_pipe[1]); - continue; - } - - pid_t pid; - { - std::unique_lock lock(snapshot_mutex_); - LOG(INFO) << "[Snapshot] Locking snapshot mutex, snapshot_id=" - << snapshot_id; - pid = fork(); - } - if (pid == -1) { - // Fork failed - LOG(ERROR) << "[Snapshot] Failed to fork child process for state " - "persistence: " - << strerror(errno) << ", snapshot_id=" << snapshot_id; - close(log_pipe[0]); - close(log_pipe[1]); - } else if (pid == 0) { - // Child process - // Close read end, set write end for logging - close(log_pipe[0]); - g_snapshot_log_pipe_fd = log_pipe[1]; - - // Save current state using the configured persistence mechanism - SNAP_LOG_INFO("[Snapshot] Child process started, snapshot_id={}", - snapshot_id); - auto result = PersistState(descriptor.value()); - if (!result) { - SNAP_LOG_ERROR( - "[Snapshot] Child process failed to persist state, " - "snapshot_id={},code={},msg={}", - snapshot_id, toString(result.error().code), - result.error().message); - close(log_pipe[1]); - _exit(1); // Exit child process with error - } - SNAP_LOG_INFO( - "[Snapshot] Child process successfully persisted state, " - "snapshot_id={}", - snapshot_id); - - close(log_pipe[1]); - _exit(0); // Exit child process successfully - } else { - // Parent process - // Close write end, pass read end to wait function - close(log_pipe[1]); - WaitForSnapshotChild(pid, snapshot_id, log_pipe[0]); - close(log_pipe[0]); - } - } - LOG(INFO) << "[Snapshot] snapshot_thread stopped"; -} - -void MasterService::WaitForSnapshotChild(pid_t pid, - const std::string& snapshot_id, - int log_pipe_fd) { - // Default 5 minute timeout - const int64_t timeout_seconds = snapshot_child_timeout_seconds_; - - LOG(INFO) - << "[Snapshot] waiting for child process to complete, snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", timeout=" << timeout_seconds << "s"; - - // Set pipe to non-blocking mode - int flags = fcntl(log_pipe_fd, F_GETFL, 0); - if (flags == -1 || fcntl(log_pipe_fd, F_SETFL, flags | O_NONBLOCK) == -1) { - LOG(WARNING) << "[Snapshot] Failed to set pipe non-blocking: " - << strerror(errno); - } - - // Buffer for reading child logs - char buf[4096]; - std::string log_buffer; - - // Helper lambda to read and output child logs - auto flush_child_logs = [&]() { - while (true) { - ssize_t n = read(log_pipe_fd, buf, sizeof(buf) - 1); - if (n > 0) { - buf[n] = '\0'; - log_buffer += buf; - // Output complete lines - size_t pos; - while ((pos = log_buffer.find('\n')) != std::string::npos) { - std::string line = log_buffer.substr(0, pos); - log_buffer.erase(0, pos + 1); - if (!line.empty()) { - LOG(INFO) << "[Snapshot:Child] " << line; - } - } - } else { - break; - } - } - }; - - // Record start time - auto start_time = std::chrono::steady_clock::now(); - - // Use non-blocking polling to wait - while (true) { - // Read child logs first - flush_child_logs(); - - int status; - pid_t result = waitpid(pid, &status, WNOHANG); - - if (result == -1) { - LOG(ERROR) << "[Snapshot] Failed to wait for child process: " - << strerror(errno) << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - return; - } else if (result == 0) { - // Child process is still running - auto elapsed = std::chrono::duration_cast( - std::chrono::steady_clock::now() - start_time) - .count(); - - if (elapsed >= timeout_seconds) { - // Timeout handling - flush remaining logs before killing - flush_child_logs(); - if (!log_buffer.empty()) { - LOG(INFO) << "[Snapshot:Child] " << log_buffer; - } - HandleChildTimeout(pid, snapshot_id); - MasterMetricManager::instance().inc_snapshot_fail(); - return; - } - - // Brief sleep before checking again - std::this_thread::sleep_for(std::chrono::seconds(2)); - } else { - // Child process has exited - // Flush remaining logs from child - flush_child_logs(); - // Output any remaining incomplete line - if (!log_buffer.empty()) { - LOG(INFO) << "[Snapshot:Child] " << log_buffer; - } - - HandleChildExit(pid, status, snapshot_id); - auto elapsed = - std::chrono::duration_cast( - std::chrono::steady_clock::now() - start_time) - .count(); - MasterMetricManager::instance().set_snapshot_duration_ms(elapsed); - return; - } - } -} - -void MasterService::HandleChildTimeout(pid_t pid, - const std::string& snapshot_id) { - LOG(WARNING) << "[Snapshot] Child process timeout, snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", killing child process"; - - // Try to gracefully terminate the child process - if (kill(pid, SIGTERM) == 0) { - // Wait a few seconds to see if it exits gracefully - std::this_thread::sleep_for(std::chrono::seconds(5)); - - // Check if it has exited - int status; - if (waitpid(pid, &status, WNOHANG) == 0) { - // Child process still not exited, force kill - LOG(WARNING) << "[Snapshot] Child process still running, force " - "killing, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - kill(pid, SIGKILL); - - // Wait for force termination to complete - waitpid(pid, &status, 0); - LOG(WARNING) - << "[Snapshot] Child process force killed, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - } else { - LOG(INFO) << "[Snapshot] Child process terminated gracefully after " - "SIGTERM, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - } - } else { - LOG(ERROR) << "[Snapshot] Failed to send SIGTERM to child process, " - "snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", error=" << strerror(errno); - } -} - -void MasterService::HandleChildExit(pid_t pid, int status, - const std::string& snapshot_id) { - if (WIFEXITED(status)) { - int exit_code = WEXITSTATUS(status); - if (exit_code != 0) { - LOG(ERROR) << "[Snapshot] Child process exited with error code: " - << exit_code << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - } else { - LOG(INFO) << "[Snapshot] Child process successfully persisted " - "state, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_success(); - } - } else if (WIFSIGNALED(status)) { - int signal = WTERMSIG(status); - LOG(ERROR) << "[Snapshot] Child process terminated by signal: " - << signal << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - } -} - -tl::expected -MasterService::ResolveSnapshotSequenceId() const { - if (!enable_ha_ || ha_backend_type_ != "etcd") { - // OpLog sequence ids start at 1. Returning 0 here is a sentinel that - // means "no persisted OpLog boundary", so a standby that later calls - // Recover(0) will replay from the first entry when oplog following is - // enabled. - return ha::OpLogSequenceId{0}; - } - -#ifndef STORE_USE_ETCD - return tl::make_unexpected(SerializationError( - ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, - "etcd snapshot sequence resolution is unavailable in this build")); -#else - auto oplog_store = GetSnapshotBoundaryOpLogStore(); - if (!oplog_store) { - return tl::make_unexpected(oplog_store.error()); - } - - uint64_t sequence_id = 0; - auto err = oplog_store.value()->GetLatestSequenceId(sequence_id); - if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - return ha::OpLogSequenceId{0}; - } - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to resolve snapshot sequence boundary: {}", - toString(err)))); - } - - return static_cast(sequence_id); -#endif -} - -#ifdef STORE_USE_ETCD -tl::expected -MasterService::GetSnapshotBoundaryOpLogStore() const { - if (ha_backend_connstring_.empty()) { - return tl::make_unexpected(SerializationError( - ErrorCode::INVALID_PARAMS, - "etcd snapshot sequence resolution requires a backend connstring")); - } - - std::lock_guard lock(snapshot_boundary_oplog_store_mutex_); - if (snapshot_boundary_oplog_store_ != nullptr) { - return snapshot_boundary_oplog_store_.get(); - } - - auto err = - EtcdHelper::ConnectToEtcdStoreClient(ha_backend_connstring_.c_str()); - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to connect to etcd for snapshot boundary: " - "{}", - toString(err)))); - } - - auto oplog_store = std::make_unique(cluster_id_); - err = oplog_store->Init(); - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to initialize etcd oplog store: {}", - toString(err)))); - } - - snapshot_boundary_oplog_store_ = std::move(oplog_store); - return snapshot_boundary_oplog_store_.get(); -} -#endif - -tl::expected -MasterService::BuildSnapshotDescriptor(const std::string& snapshot_id, - const std::string& manifest_path, - const std::string& object_prefix) const { - auto sequence_id = ResolveSnapshotSequenceId(); - if (!sequence_id) { - return tl::make_unexpected(sequence_id.error()); - } - - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor( - snapshot_root, snapshot_id); - descriptor.last_included_seq = sequence_id.value(); - descriptor.producer_view_version = view_version_; - descriptor.manifest_key = manifest_path; - descriptor.object_prefix = object_prefix; - descriptor.created_at_ms = CurrentTimeMs(); - return descriptor; -} - -tl::expected MasterService::PersistState( - const std::string& snapshot_id) { - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - const std::string path_prefix = snapshot_root + snapshot_id + "/"; - const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - auto descriptor = - BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); - if (!descriptor) { - return tl::make_unexpected(descriptor.error()); - } - return PersistState(descriptor.value()); -} - -tl::expected MasterService::PersistState( - const ha::SnapshotDescriptor& descriptor) { - const std::string& snapshot_id = descriptor.snapshot_id; - const std::string& path_prefix = descriptor.object_prefix; - const std::string& manifest_path = descriptor.manifest_key; - - try { - auto* snapshot_catalog_store = GetSnapshotCatalogStore(); - if (!snapshot_catalog_store) { - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - "snapshot catalog store is not initialized")); - } - - SNAP_LOG_INFO( - "[Snapshot] action=persisting_state start, snapshot_id={}, " - "serializer_type={}, version={}", - snapshot_id, SNAPSHOT_SERIALIZER_TYPE, SNAPSHOT_SERIALIZER_VERSION); - MetadataSerializer metadata_serializer(this); - SegmentSerializer segment_serializer(&segment_manager_); - TaskManagerSerializer task_manager_serializer(&task_manager_); - - auto metadata_result = metadata_serializer.Serialize(); - if (!metadata_result) { - SNAP_LOG_ERROR( - "[Snapshot] metadata serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(metadata_result.error().code), - metadata_result.error().message); - - return tl::make_unexpected(metadata_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] metadata serialization_successful, snapshot_id={}", - snapshot_id); - - auto segment_result = segment_serializer.Serialize(); - if (!segment_result) { - SNAP_LOG_ERROR( - "[Snapshot] segment serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(segment_result.error().code), - segment_result.error().message); - return tl::make_unexpected(segment_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] segment serialization_successful, snapshot_id={}", - snapshot_id); - - auto task_manager_result = task_manager_serializer.Serialize(); - if (!task_manager_result) { - SNAP_LOG_ERROR( - "[Snapshot] task manager serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(task_manager_result.error().code), - task_manager_result.error().message); - return tl::make_unexpected(task_manager_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] task manager serialization_successful, snapshot_id={}", - snapshot_id); - - const auto& serialized_metadata = metadata_result.value(); - const auto& serialized_segment = segment_result.value(); - const auto& serialized_task_manager = task_manager_result.value(); - - // When backup_dir is enabled, try all uploads to ensure complete backup - // When backup_dir is disabled, use fail-fast mode - bool upload_success = true; - std::string error_msg; - SNAP_LOG_INFO("[Snapshot] Backend info: {}", - snapshot_object_store_->GetConnectionInfo()); - - // Upload metadata - std::string metadata_path = path_prefix + SNAPSHOT_METADATA_FILE; - auto upload_result = - UploadSnapshotPayloadFile(serialized_metadata, metadata_path, - SNAPSHOT_METADATA_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] metadata upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, metadata_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload segment - std::string segment_path = path_prefix + SNAPSHOT_SEGMENTS_FILE; - upload_result = - UploadSnapshotPayloadFile(serialized_segment, segment_path, - SNAPSHOT_SEGMENTS_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] segment upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, segment_path, toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload task manager - std::string task_manager_path = - path_prefix + SNAPSHOT_TASK_MANAGER_FILE; - upload_result = UploadSnapshotPayloadFile( - serialized_task_manager, task_manager_path, - SNAPSHOT_TASK_MANAGER_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] task_manager upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, task_manager_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload manifest - std::string manifest_content = - fmt::format("{}|{}|{}", SNAPSHOT_SERIALIZER_TYPE, - SNAPSHOT_SERIALIZER_VERSION, snapshot_id); - std::vector manifest_bytes(manifest_content.begin(), - manifest_content.end()); - upload_result = UploadSnapshotPayloadFile( - manifest_bytes, manifest_path, SNAPSHOT_MANIFEST_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] manifest upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, manifest_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - if (!upload_success) { - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); - } - - // Publish snapshot catalog entry and advance the latest marker. - std::string latest_path = - snapshot_catalog_store->GetSnapshotRoot() + SNAPSHOT_LATEST_FILE; - std::string latest_content = snapshot_id; - - auto publish_result = snapshot_catalog_store->Publish(descriptor); - if (publish_result != ErrorCode::OK) { - SNAP_LOG_ERROR( - "[Snapshot] latest update failed, snapshot_id={}, file={}, " - "code={}", - snapshot_id, latest_path, toString(publish_result)); - if (use_snapshot_backup_dir_) { - auto save_path = fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_SAVE_DIR / - SNAPSHOT_LATEST_FILE; - auto save_result = - FileUtil::SaveStringToFile(latest_content, save_path); - if (!save_result) { - SNAP_LOG_ERROR( - "[Snapshot] save latest to disk failed, " - "snapshot_id={}, " - "content={}, file={}", - snapshot_id, latest_content, save_path.string()); - } - } - - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - fmt::format("latest update {} failed", latest_path))); - } - SNAP_LOG_INFO( - "[Snapshot] Upload latest success: {}, snapshot_id={}, " - "content={}", - latest_path, snapshot_id, latest_content); - - CleanupOldSnapshot(snapshot_retention_count_, snapshot_id); - SNAP_LOG_INFO("[Snapshot] action=persisting_state end, snapshot_id={}", - snapshot_id); - } catch (const std::exception& e) { - SNAP_LOG_ERROR( - "[Snapshot] Exception during state persistent, snapshot_id={}, " - "error={}", - snapshot_id, e.what()); - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - fmt::format("Exception during state persistent: {}", e.what()))); - } catch (...) { - SNAP_LOG_ERROR( - "[Snapshot] Unknown exception during state persistent, " - "snapshot_id={}", - snapshot_id); - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, - "Unknown exception during state persistent")); - } - return {}; -} - -tl::expected MasterService::UploadSnapshotPayloadFile( - const std::vector& data, const std::string& path, - const std::string& local_filename, const std::string& snapshot_id) { - SNAP_LOG_INFO("[Snapshot] Uploading {} to: {}, snapshot_id={}", - local_filename, path, snapshot_id); - - std::string error_msg; - auto upload_result = snapshot_object_store_->UploadBuffer(path, data); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] {} upload failed, snapshot_id={}, file={}, error={}", - local_filename, snapshot_id, path, upload_result.error()); - - // Upload failed, save locally for manual recovery in exception - // scenarios - if (use_snapshot_backup_dir_) { - auto save_path = fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_SAVE_DIR / local_filename; - auto save_result = FileUtil::SaveBinaryToFile(data, save_path); - if (!save_result) { - SNAP_LOG_ERROR( - "[Snapshot] save {} to disk failed, snapshot_id={}, " - "file={}", - local_filename, snapshot_id, save_path.string()); - } - } - - error_msg.append(local_filename) - .append(" upload ") - .append(path) - .append(" failed; "); - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); - } else { - SNAP_LOG_INFO("[Snapshot] Upload {} success: {}, snapshot_id={}", - local_filename, path, snapshot_id); - } - - return {}; -} - -void MasterService::CleanupOldSnapshot(int keep_count, - const std::string& snapshot_id) { - auto* snapshot_catalog_store = GetSnapshotCatalogStore(); - if (!snapshot_catalog_store) { - SNAP_LOG_ERROR( - "[Snapshot] snapshot catalog store is not initialized, " - "snapshot_id={}", - snapshot_id); - return; - } - - // List() loads one descriptor per published snapshot. This remains cheap - // because CleanupOldSnapshot() itself enforces snapshot_retention_count_ - // and keeps the catalog single-digit in normal deployments. - auto list_result = snapshot_catalog_store->List(kUnlimitedSnapshotList); - if (!list_result) { - SNAP_LOG_ERROR("[Snapshot] error=list failed, snapshot_id={}, code={}", - snapshot_id, toString(list_result.error())); - return; - } - - const auto& snapshots = list_result.value(); - - if (static_cast(snapshots.size()) > keep_count) { - for (int i = keep_count; i < static_cast(snapshots.size()); i++) { - const std::string& old_state_dir = snapshots[i].snapshot_id; - - if (old_state_dir == snapshot_id) { - SNAP_LOG_WARN( - "[Snapshot] Skipping deletion of current snapshot " - "directory {}, " - "snapshot_id={}", - old_state_dir, snapshot_id); - continue; - } - - auto delete_result = snapshot_catalog_store->Delete(old_state_dir); - if (delete_result != ErrorCode::OK) { - SNAP_LOG_ERROR( - "[Snapshot] Failed to delete old snapshot {}, " - "snapshot_id={}, code={}", - old_state_dir, snapshot_id, toString(delete_result)); - } else { - SNAP_LOG_INFO( - "[Snapshot] Successfully deleted old snapshot {}, " - "snapshot_id={}", - old_state_dir, snapshot_id); - } - } - } -} - void MasterService::RestoreState() { auto* snapshot_catalog_store = GetSnapshotCatalogStore(); if (!snapshot_catalog_store) { @@ -8766,23 +8112,6 @@ MasterService::MetadataSerializer::DeserializeMetadata( return metadata; } -std::string MasterService::FormatTimestamp( - const std::chrono::system_clock::time_point& tp) { - auto time_t = std::chrono::system_clock::to_time_t(tp); - - std::stringstream ss; - ss << std::put_time(std::localtime(&time_t), "%Y%m%d_%H%M%S"); - - // Add milliseconds to ensure uniqueness - auto ms = std::chrono::duration_cast( - tp.time_since_epoch()) % - 1000; - - ss << "_" << std::setfill('0') << std::setw(3) << ms.count(); - - return ss.str(); -} - tl::expected MasterService::CreateCopyTask( const std::string& key, const std::string& tenant_id, const std::vector& targets) { diff --git a/mooncake-store/src/master_snapshot_manager.cpp b/mooncake-store/src/master_snapshot_manager.cpp new file mode 100644 index 0000000000..13980b57b3 --- /dev/null +++ b/mooncake-store/src/master_snapshot_manager.cpp @@ -0,0 +1,712 @@ +#include "master_snapshot_manager.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "master_service.h" +#include "master_metric_manager.h" +#include "master_snapshot_repository.h" +#include "ha/snapshot/catalog/snapshot_catalog_store.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "ha/snapshot/snapshot_logger.h" +#include "serialize/serializer.h" +#include "segment.h" +#include "task_manager.h" +#include "utils/file_util.h" +#include "utils/zstd_util.h" + +#ifdef STORE_USE_ETCD +#include "ha/oplog/etcd_oplog_store.h" +#include "etcd_helper.h" +#endif + +namespace mooncake { + +// Snapshot file names (moved from master_service.cpp) +static const std::string SNAPSHOT_METADATA_FILE = "metadata"; +static const std::string SNAPSHOT_SEGMENTS_FILE = "segments"; +static const std::string SNAPSHOT_TASK_MANAGER_FILE = "task_manager"; +static const std::string SNAPSHOT_MANIFEST_FILE = "manifest.txt"; +static const std::string SNAPSHOT_LATEST_FILE = "latest.txt"; +static const std::string SNAPSHOT_BACKUP_SAVE_DIR = + "mooncake_snapshot_save_backup"; +static const std::string SNAPSHOT_SERIALIZER_VERSION = "1.0.0"; +static const std::string SNAPSHOT_SERIALIZER_TYPE = "messagepack"; + +namespace { +int64_t CurrentTimeMs() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} +} // namespace + +MasterSnapshotManager::MasterSnapshotManager( + MasterService* master_service, MasterSnapshotManagerOptions options, + std::shared_mutex& snapshot_mutex, + SnapshotObjectStore* snapshot_object_store, + ha::SnapshotCatalogStore* snapshot_catalog_store) + : master_service_(master_service), + options_(std::move(options)), + snapshot_mutex_(snapshot_mutex), + snapshot_object_store_(snapshot_object_store), + snapshot_catalog_store_(snapshot_catalog_store), + repository_(std::make_unique( + snapshot_object_store, snapshot_catalog_store, + options_.snapshot_backup_dir, options_.use_snapshot_backup_dir)) {} + +MasterSnapshotManager::~MasterSnapshotManager() { + Stop(); + if (snapshot_thread_.joinable()) { + snapshot_thread_.join(); + } +} + +void MasterSnapshotManager::Start() { + if (snapshot_running_.load()) { + return; + } + snapshot_running_ = true; + snapshot_thread_ = + std::thread(&MasterSnapshotManager::SnapshotThreadFunc, this); + LOG(INFO) << "[MasterSnapshotManager] Started"; +} + +void MasterSnapshotManager::Stop() { + { + std::lock_guard lk(snapshot_thread_mutex_); + snapshot_running_ = false; + } + snapshot_thread_cv_.notify_all(); + LOG(INFO) << "[MasterSnapshotManager] Stop signaled"; +} + +std::string MasterSnapshotManager::FormatTimestamp( + const std::chrono::system_clock::time_point& tp) { + auto time_t = std::chrono::system_clock::to_time_t(tp); + + std::stringstream ss; + std::tm tm_now; + localtime_r(&time_t, &tm_now); + ss << std::put_time(&tm_now, "%Y%m%d_%H%M%S"); + + // Add milliseconds to ensure uniqueness + auto ms = std::chrono::duration_cast( + tp.time_since_epoch()) % + 1000; + + ss << "_" << std::setfill('0') << std::setw(3) << ms.count(); + + return ss.str(); +} + +void MasterSnapshotManager::SnapshotThreadFunc() { + LOG(INFO) << "[Snapshot] snapshot_thread started"; + while (snapshot_running_) { + // Wait for the next snapshot cycle, but allow fast shutdown. + { + std::unique_lock lk(snapshot_thread_mutex_); + snapshot_thread_cv_.wait_for( + lk, std::chrono::seconds(options_.snapshot_interval_seconds), + [&] { return !snapshot_running_.load(); }); + } + + if (!snapshot_running_) { + break; + } + + if (!options_.enable_snapshot) { + // Snapshot is disabled + LOG(INFO) + << "[Snapshot] Snapshot is disabled, waiting for next cycle"; + continue; + } + // Fork a child process to save current state + + std::string snapshot_id = + FormatTimestamp(std::chrono::system_clock::now()); + LOG(INFO) << "[Snapshot] Preparing to fork child process, snapshot_id=" + << snapshot_id; + + // Create pipe for child process logging + int log_pipe[2]; + if (pipe(log_pipe) == -1) { + LOG(ERROR) << "[Snapshot] Failed to create log pipe: " + << strerror(errno) << ", snapshot_id=" << snapshot_id; + continue; + } + + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + const std::string path_prefix = snapshot_root + snapshot_id + "/"; + const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; + auto descriptor = + BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); + if (!descriptor) { + LOG(ERROR) << "[Snapshot] Failed to build descriptor before fork, " + "snapshot_id=" + << snapshot_id + << ", code=" << toString(descriptor.error().code) + << ", msg=" << descriptor.error().message; + close(log_pipe[0]); + close(log_pipe[1]); + continue; + } + + pid_t pid; + { + std::unique_lock lock(snapshot_mutex_); + LOG(INFO) << "[Snapshot] Locking snapshot mutex, snapshot_id=" + << snapshot_id; + pid = fork(); + } + if (pid == -1) { + // Fork failed + LOG(ERROR) << "[Snapshot] Failed to fork child process for state " + "persistence: " + << strerror(errno) << ", snapshot_id=" << snapshot_id; + close(log_pipe[0]); + close(log_pipe[1]); + } else if (pid == 0) { + // Child process + // Close read end, set write end for logging + close(log_pipe[0]); + g_snapshot_log_pipe_fd = log_pipe[1]; + + // Save current state using the configured persistence mechanism + SNAP_LOG_INFO("[Snapshot] Child process started, snapshot_id={}", + snapshot_id); + auto result = PersistState(descriptor.value()); + if (!result) { + SNAP_LOG_ERROR( + "[Snapshot] Child process failed to persist state, " + "snapshot_id={},code={},msg={}", + snapshot_id, toString(result.error().code), + result.error().message); + close(log_pipe[1]); + _exit(1); // Exit child process with error + } + SNAP_LOG_INFO( + "[Snapshot] Child process successfully persisted state, " + "snapshot_id={}", + snapshot_id); + + close(log_pipe[1]); + _exit(0); // Exit child process successfully + } else { + // Parent process + // Close write end, pass read end to wait function + close(log_pipe[1]); + WaitForSnapshotChild(pid, snapshot_id, log_pipe[0]); + close(log_pipe[0]); + } + } + LOG(INFO) << "[Snapshot] snapshot_thread stopped"; +} + +void MasterSnapshotManager::WaitForSnapshotChild(pid_t pid, + const std::string& snapshot_id, + int log_pipe_fd) { + // Default 5 minute timeout + const int64_t timeout_seconds = options_.snapshot_child_timeout_seconds; + + LOG(INFO) + << "[Snapshot] waiting for child process to complete, snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", timeout=" << timeout_seconds << "s"; + + // Set pipe to non-blocking mode + int flags = fcntl(log_pipe_fd, F_GETFL, 0); + if (flags == -1 || fcntl(log_pipe_fd, F_SETFL, flags | O_NONBLOCK) == -1) { + LOG(WARNING) << "[Snapshot] Failed to set pipe non-blocking: " + << strerror(errno); + } + + // Buffer for reading child logs + char buf[4096]; + std::string log_buffer; + + // Helper lambda to read and output child logs + auto flush_child_logs = [&]() { + while (true) { + ssize_t n = read(log_pipe_fd, buf, sizeof(buf) - 1); + if (n > 0) { + buf[n] = '\0'; + log_buffer += buf; + // Output complete lines + size_t pos; + while ((pos = log_buffer.find('\n')) != std::string::npos) { + std::string line = log_buffer.substr(0, pos); + log_buffer.erase(0, pos + 1); + if (!line.empty()) { + LOG(INFO) << "[Snapshot:Child] " << line; + } + } + } else { + break; + } + } + }; + + // Record start time + auto start_time = std::chrono::steady_clock::now(); + + // Use non-blocking polling to wait + while (true) { + // Read child logs first + flush_child_logs(); + + int status; + pid_t result = waitpid(pid, &status, WNOHANG); + + if (result == -1) { + LOG(ERROR) << "[Snapshot] Failed to wait for child process: " + << strerror(errno) << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + return; + } else if (result == 0) { + // Child process is still running + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time) + .count(); + + if (elapsed >= timeout_seconds) { + // Timeout handling - flush remaining logs before killing + flush_child_logs(); + if (!log_buffer.empty()) { + LOG(INFO) << "[Snapshot:Child] " << log_buffer; + } + HandleChildTimeout(pid, snapshot_id); + MasterMetricManager::instance().inc_snapshot_fail(); + return; + } + + // Brief sleep before checking again + std::this_thread::sleep_for(std::chrono::seconds(2)); + } else { + // Child process has exited + // Flush remaining logs from child + flush_child_logs(); + // Output any remaining incomplete line + if (!log_buffer.empty()) { + LOG(INFO) << "[Snapshot:Child] " << log_buffer; + } + + HandleChildExit(pid, status, snapshot_id); + auto elapsed = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time) + .count(); + MasterMetricManager::instance().set_snapshot_duration_ms(elapsed); + return; + } + } +} + +void MasterSnapshotManager::HandleChildTimeout(pid_t pid, + const std::string& snapshot_id) { + LOG(WARNING) << "[Snapshot] Child process timeout, snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", killing child process"; + + // Try to gracefully terminate the child process + if (kill(pid, SIGTERM) == 0) { + // Wait a few seconds to see if it exits gracefully + std::this_thread::sleep_for(std::chrono::seconds(5)); + + // Check if it has exited + int status; + if (waitpid(pid, &status, WNOHANG) == 0) { + // Child process still not exited, force kill + LOG(WARNING) << "[Snapshot] Child process still running, force " + "killing, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + kill(pid, SIGKILL); + + // Wait for force termination to complete + waitpid(pid, &status, 0); + LOG(WARNING) + << "[Snapshot] Child process force killed, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + } else { + LOG(INFO) << "[Snapshot] Child process terminated gracefully after " + "SIGTERM, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + } + } else { + LOG(ERROR) << "[Snapshot] Failed to send SIGTERM to child process, " + "snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", error=" << strerror(errno); + } +} + +void MasterSnapshotManager::HandleChildExit(pid_t pid, int status, + const std::string& snapshot_id) { + if (WIFEXITED(status)) { + int exit_code = WEXITSTATUS(status); + if (exit_code != 0) { + LOG(ERROR) << "[Snapshot] Child process exited with error code: " + << exit_code << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + } else { + LOG(INFO) << "[Snapshot] Child process successfully persisted " + "state, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_success(); + } + } else if (WIFSIGNALED(status)) { + int signal = WTERMSIG(status); + LOG(ERROR) << "[Snapshot] Child process terminated by signal: " + << signal << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + } +} + +tl::expected +MasterSnapshotManager::ResolveSnapshotSequenceId() const { + if (!options_.enable_ha || options_.ha_backend_type != "etcd") { + // OpLog sequence ids start at 1. Returning 0 here is a sentinel that + // means "no persisted OpLog boundary", so a standby that later calls + // Recover(0) will replay from the first entry when oplog following is + // enabled. + return ha::OpLogSequenceId{0}; + } + +#ifndef STORE_USE_ETCD + return tl::make_unexpected(SerializationError( + ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, + "etcd snapshot sequence resolution is unavailable in this build")); +#else + auto oplog_store = GetSnapshotBoundaryOpLogStore(); + if (!oplog_store) { + return tl::make_unexpected(oplog_store.error()); + } + + uint64_t sequence_id = 0; + auto err = oplog_store.value()->GetLatestSequenceId(sequence_id); + if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { + return ha::OpLogSequenceId{0}; + } + if (err != ErrorCode::OK) { + return tl::make_unexpected(SerializationError( + err, fmt::format("failed to resolve snapshot sequence boundary: {}", + toString(err)))); + } + + return static_cast(sequence_id); +#endif +} + +#ifdef STORE_USE_ETCD +tl::expected +MasterSnapshotManager::GetSnapshotBoundaryOpLogStore() const { + if (options_.ha_backend_connstring.empty()) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, + "etcd snapshot sequence resolution requires a backend connstring")); + } + + std::lock_guard lock(snapshot_boundary_oplog_store_mutex_); + if (snapshot_boundary_oplog_store_ != nullptr) { + return snapshot_boundary_oplog_store_.get(); + } + + auto err = EtcdHelper::ConnectToEtcdStoreClient( + options_.ha_backend_connstring.c_str()); + if (err != ErrorCode::OK) { + return tl::make_unexpected(SerializationError( + err, fmt::format("failed to connect to etcd for snapshot boundary: " + "{}", + toString(err)))); + } + + auto oplog_store = std::make_unique(options_.cluster_id); + err = oplog_store->Init(); + if (err != ErrorCode::OK) { + return tl::make_unexpected(SerializationError( + err, fmt::format("failed to initialize etcd oplog store: {}", + toString(err)))); + } + + snapshot_boundary_oplog_store_ = std::move(oplog_store); + return snapshot_boundary_oplog_store_.get(); +} +#endif + +tl::expected +MasterSnapshotManager::BuildSnapshotDescriptor( + const std::string& snapshot_id, const std::string& manifest_path, + const std::string& object_prefix) const { + auto sequence_id = ResolveSnapshotSequenceId(); + if (!sequence_id) { + return tl::make_unexpected(sequence_id.error()); + } + + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor( + snapshot_root, snapshot_id); + descriptor.last_included_seq = sequence_id.value(); + descriptor.producer_view_version = master_service_->view_version_; + descriptor.manifest_key = manifest_path; + descriptor.object_prefix = object_prefix; + descriptor.created_at_ms = CurrentTimeMs(); + return descriptor; +} + +tl::expected MasterSnapshotManager::PersistState( + const std::string& snapshot_id) { + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + const std::string path_prefix = snapshot_root + snapshot_id + "/"; + const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; + auto descriptor = + BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); + if (!descriptor) { + return tl::make_unexpected(descriptor.error()); + } + return PersistState(descriptor.value()); +} + +tl::expected MasterSnapshotManager::PersistState( + const ha::SnapshotDescriptor& descriptor) { + const std::string& snapshot_id = descriptor.snapshot_id; + const std::string& path_prefix = descriptor.object_prefix; + const std::string& manifest_path = descriptor.manifest_key; + + try { + if (!snapshot_catalog_store_) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + "snapshot catalog store is not initialized")); + } + + SNAP_LOG_INFO( + "[Snapshot] action=persisting_state start, snapshot_id={}, " + "serializer_type={}, version={}", + snapshot_id, SNAPSHOT_SERIALIZER_TYPE, SNAPSHOT_SERIALIZER_VERSION); + MasterService::MetadataSerializer metadata_serializer(master_service_); + SegmentSerializer segment_serializer( + &master_service_->segment_manager_); + TaskManagerSerializer task_manager_serializer( + &master_service_->task_manager_); + + auto metadata_result = metadata_serializer.Serialize(); + if (!metadata_result) { + SNAP_LOG_ERROR( + "[Snapshot] metadata serialization failed, snapshot_id={}, " + "code={}, msg={}", + snapshot_id, toString(metadata_result.error().code), + metadata_result.error().message); + + return tl::make_unexpected(metadata_result.error()); + } + SNAP_LOG_INFO( + "[Snapshot] metadata serialization_successful, snapshot_id={}", + snapshot_id); + + auto segment_result = segment_serializer.Serialize(); + if (!segment_result) { + SNAP_LOG_ERROR( + "[Snapshot] segment serialization failed, snapshot_id={}, " + "code={}, msg={}", + snapshot_id, toString(segment_result.error().code), + segment_result.error().message); + return tl::make_unexpected(segment_result.error()); + } + SNAP_LOG_INFO( + "[Snapshot] segment serialization_successful, snapshot_id={}", + snapshot_id); + + auto task_manager_result = task_manager_serializer.Serialize(); + if (!task_manager_result) { + SNAP_LOG_ERROR( + "[Snapshot] task manager serialization failed, snapshot_id={}, " + "code={}, msg={}", + snapshot_id, toString(task_manager_result.error().code), + task_manager_result.error().message); + return tl::make_unexpected(task_manager_result.error()); + } + SNAP_LOG_INFO( + "[Snapshot] task manager serialization_successful, snapshot_id={}", + snapshot_id); + + const auto& serialized_metadata = metadata_result.value(); + const auto& serialized_segment = segment_result.value(); + const auto& serialized_task_manager = task_manager_result.value(); + + // When backup_dir is enabled, try all uploads to ensure complete backup + // When backup_dir is disabled, use fail-fast mode + bool upload_success = true; + std::string error_msg; + SNAP_LOG_INFO("[Snapshot] Backend info: {}", + repository_->GetObjectStoreConnectionInfo()); + + // Upload metadata + std::string metadata_path = path_prefix + SNAPSHOT_METADATA_FILE; + auto upload_result = + repository_->UploadPayloadFile(serialized_metadata, metadata_path, + SNAPSHOT_METADATA_FILE, snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] metadata upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, metadata_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload segment + std::string segment_path = path_prefix + SNAPSHOT_SEGMENTS_FILE; + upload_result = + repository_->UploadPayloadFile(serialized_segment, segment_path, + SNAPSHOT_SEGMENTS_FILE, snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] segment upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, segment_path, toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload task manager + std::string task_manager_path = + path_prefix + SNAPSHOT_TASK_MANAGER_FILE; + upload_result = repository_->UploadPayloadFile( + serialized_task_manager, task_manager_path, + SNAPSHOT_TASK_MANAGER_FILE, snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] task_manager upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, task_manager_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload manifest + std::string manifest_content = + fmt::format("{}|{}|{}", SNAPSHOT_SERIALIZER_TYPE, + SNAPSHOT_SERIALIZER_VERSION, snapshot_id); + std::vector manifest_bytes(manifest_content.begin(), + manifest_content.end()); + upload_result = repository_->UploadPayloadFile( + manifest_bytes, manifest_path, SNAPSHOT_MANIFEST_FILE, snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] manifest upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, manifest_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + if (!upload_success) { + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); + } + + // Publish snapshot catalog entry and advance the latest marker. + std::string latest_path = + snapshot_catalog_store_->GetSnapshotRoot() + SNAPSHOT_LATEST_FILE; + std::string latest_content = snapshot_id; + + auto publish_result = repository_->PublishSnapshot(descriptor); + if (publish_result != ErrorCode::OK) { + SNAP_LOG_ERROR( + "[Snapshot] latest update failed, snapshot_id={}, file={}, " + "code={}", + snapshot_id, latest_path, toString(publish_result)); + if (options_.use_snapshot_backup_dir) { + auto save_path = fs::path(options_.snapshot_backup_dir) / + SNAPSHOT_BACKUP_SAVE_DIR / + SNAPSHOT_LATEST_FILE; + auto save_result = + FileUtil::SaveStringToFile(latest_content, save_path); + if (!save_result) { + SNAP_LOG_ERROR( + "[Snapshot] save latest to disk failed, " + "snapshot_id={}, " + "content={}, file={}", + snapshot_id, latest_content, save_path.string()); + } + } + + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + fmt::format("latest update {} failed", latest_path))); + } + SNAP_LOG_INFO( + "[Snapshot] Upload latest success: {}, snapshot_id={}, " + "content={}", + latest_path, snapshot_id, latest_content); + + repository_->CleanupOldSnapshots(options_.snapshot_retention_count, + snapshot_id); + SNAP_LOG_INFO("[Snapshot] action=persisting_state end, snapshot_id={}", + snapshot_id); + } catch (const std::exception& e) { + SNAP_LOG_ERROR( + "[Snapshot] Exception during state persistent, snapshot_id={}, " + "error={}", + snapshot_id, e.what()); + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + fmt::format("Exception during state persistent: {}", e.what()))); + } catch (...) { + SNAP_LOG_ERROR( + "[Snapshot] Unknown exception during state persistent, " + "snapshot_id={}", + snapshot_id); + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, + "Unknown exception during state persistent")); + } + return {}; +} + +tl::expected +MasterSnapshotManager::UploadSnapshotPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id) { + return repository_->UploadPayloadFile(data, path, local_filename, + snapshot_id); +} + +void MasterSnapshotManager::CleanupOldSnapshot(size_t keep_count, + const std::string& snapshot_id) { + repository_->CleanupOldSnapshots(keep_count, snapshot_id); +} + +} // namespace mooncake diff --git a/mooncake-store/src/master_snapshot_repository.cpp b/mooncake-store/src/master_snapshot_repository.cpp new file mode 100644 index 0000000000..e0f63a5a06 --- /dev/null +++ b/mooncake-store/src/master_snapshot_repository.cpp @@ -0,0 +1,142 @@ +#include "master_snapshot_repository.h" + +#include + +#include "ha/snapshot/catalog/snapshot_catalog_store.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "ha/snapshot/snapshot_logger.h" +#include "utils/file_util.h" + +namespace mooncake { + +namespace fs = std::filesystem; + +namespace { +constexpr size_t kUnlimitedSnapshotList = 0; +static const std::string SNAPSHOT_BACKUP_SAVE_DIR = + "mooncake_snapshot_save_backup"; +} // namespace + +MasterSnapshotRepository::MasterSnapshotRepository( + SnapshotObjectStore* object_store, ha::SnapshotCatalogStore* catalog_store, + const std::string& backup_dir, bool use_backup_dir) + : object_store_(object_store), + catalog_store_(catalog_store), + backup_dir_(backup_dir), + use_backup_dir_(use_backup_dir) {} + +tl::expected +MasterSnapshotRepository::UploadPayloadFile(const std::vector& data, + const std::string& path, + const std::string& local_filename, + const std::string& snapshot_id) { + SNAP_LOG_INFO("[Snapshot] Uploading {} to: {}, snapshot_id={}", + local_filename, path, snapshot_id); + + std::string error_msg; + auto upload_result = object_store_->UploadBuffer(path, data); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] {} upload failed, snapshot_id={}, file={}, error={}", + local_filename, snapshot_id, path, upload_result.error()); + + // Upload failed, save locally for manual recovery in exception + // scenarios + if (use_backup_dir_) { + auto save_path = fs::path(backup_dir_) / SNAPSHOT_BACKUP_SAVE_DIR / + local_filename; + auto save_result = FileUtil::SaveBinaryToFile(data, save_path); + if (!save_result) { + SNAP_LOG_ERROR( + "[Snapshot] save {} to disk failed, snapshot_id={}, " + "file={}", + local_filename, snapshot_id, save_path.string()); + } + } + + error_msg.append(local_filename) + .append(" upload ") + .append(path) + .append(" failed; "); + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); + } else { + SNAP_LOG_INFO("[Snapshot] Upload {} success: {}, snapshot_id={}", + local_filename, path, snapshot_id); + } + + return {}; +} + +ErrorCode MasterSnapshotRepository::PublishSnapshot( + const ha::SnapshotDescriptor& descriptor) { + return catalog_store_->Publish(descriptor); +} + +void MasterSnapshotRepository::CleanupOldSnapshots( + size_t keep_count, const std::string& current_snapshot_id) { + if (!catalog_store_) { + SNAP_LOG_ERROR( + "[Snapshot] snapshot catalog store is not initialized, " + "snapshot_id={}", + current_snapshot_id); + return; + } + + // List() loads one descriptor per published snapshot. This remains cheap + // because CleanupOldSnapshots() itself enforces retention count + // and keeps the catalog single-digit in normal deployments. + auto list_result = catalog_store_->List(kUnlimitedSnapshotList); + if (!list_result) { + SNAP_LOG_ERROR("[Snapshot] error=list failed, snapshot_id={}, code={}", + current_snapshot_id, toString(list_result.error())); + return; + } + + const auto& snapshots = list_result.value(); + + if (snapshots.size() > keep_count) { + for (size_t i = keep_count; i < snapshots.size(); i++) { + const std::string& old_state_dir = snapshots[i].snapshot_id; + + if (old_state_dir == current_snapshot_id) { + SNAP_LOG_WARN( + "[Snapshot] Skipping deletion of current snapshot " + "directory {}, " + "snapshot_id={}", + old_state_dir, current_snapshot_id); + continue; + } + + auto delete_result = catalog_store_->Delete(old_state_dir); + if (delete_result != ErrorCode::OK) { + SNAP_LOG_ERROR( + "[Snapshot] Failed to delete old snapshot {}, " + "snapshot_id={}, code={}", + old_state_dir, current_snapshot_id, + toString(delete_result)); + } else { + SNAP_LOG_INFO( + "[Snapshot] Successfully deleted old snapshot {}, " + "snapshot_id={}", + old_state_dir, current_snapshot_id); + } + } + } +} + +tl::expected, ErrorCode> +MasterSnapshotRepository::ListSnapshots(size_t limit) { + return catalog_store_->List(limit); +} + +ErrorCode MasterSnapshotRepository::DeleteSnapshot( + const ha::SnapshotId& snapshot_id) { + return catalog_store_->Delete(snapshot_id); +} + +std::string MasterSnapshotRepository::GetObjectStoreConnectionInfo() const { + return object_store_->GetConnectionInfo(); +} + +} // namespace mooncake diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h index 5d48eb25e2..6899639175 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h +++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h @@ -1,6 +1,7 @@ #pragma once #include "master_service.h" +#include "master_snapshot_manager.h" #include "master_metric_manager.h" #include "segment.h" #include "ha/snapshot/catalog/snapshot_catalog_store.h" @@ -159,11 +160,43 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { // ==================== Snapshot Helper Methods ==================== - // Wrapper method: Call MasterService's private method PersistState - // This class is a friend of MasterService, so it can access private members + // Wrapper method: Call MasterSnapshotManager's PersistState through + // MasterService This class is a friend of MasterService, so it can access + // private members static tl::expected CallPersistState( MasterService* service, const std::string& snapshot_id) { - return service->PersistState(snapshot_id); + // If snapshot_manager_ exists, use it; otherwise create a temporary one + if (service->snapshot_manager_) { + return service->snapshot_manager_->PersistState(snapshot_id); + } + + // For tests that don't have snapshot_manager_ initialized, + // we need to access the old implementation or create a temporary + // manager This is a temporary compatibility layer for tests + EnsureSnapshotStores(service); + + MasterSnapshotManagerOptions options; + options.enable_snapshot = true; + options.snapshot_interval_seconds = 300; + options.snapshot_child_timeout_seconds = 300; + options.snapshot_retention_count = 3; + options.snapshot_backup_dir = ""; + options.use_snapshot_backup_dir = false; + options.snapshot_catalog_store_type = + service->snapshot_catalog_store_type_; + options.snapshot_catalog_store_connstring = + service->snapshot_catalog_store_connstring_; + options.ha_backend_type = service->ha_backend_type_; + options.ha_backend_connstring = service->ha_backend_connstring_; + options.cluster_id = service->cluster_id_; + options.enable_ha = service->enable_ha_; + + auto temp_manager = std::make_unique( + service, options, service->snapshot_mutex_, + service->snapshot_object_store_.get(), + service->snapshot_catalog_store_.get()); + + return temp_manager->PersistState(snapshot_id); } static void EnsureSnapshotStores(MasterService* service) { diff --git a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp index 8902329a39..464db29043 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp +++ b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp @@ -1,4 +1,5 @@ #include "master_service.h" +#include "master_snapshot_manager.h" #include "master_metric_manager.h" #include "ha/snapshot/catalog/snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" @@ -115,28 +116,57 @@ class SnapshotChildProcessTest : public ::testing::Test { // Helper wrappers for private methods (friend access) std::string CallFormatTimestamp( const std::chrono::system_clock::time_point& tp) { - return service_->FormatTimestamp(tp); + // FormatTimestamp is now in MasterSnapshotManager + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->FormatTimestamp(tp); + } + // Fallback for tests without snapshot_manager_ + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->FormatTimestamp(tp); } void CallHandleChildExit(pid_t pid, int status, const std::string& snapshot_id) { - service_->HandleChildExit(pid, status, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->HandleChildExit(pid, status, + snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->HandleChildExit(pid, status, snapshot_id); + } } void CallHandleChildTimeout(pid_t pid, const std::string& snapshot_id) { - service_->HandleChildTimeout(pid, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->HandleChildTimeout(pid, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->HandleChildTimeout(pid, snapshot_id); + } } void CallCleanupOldSnapshot(int keep_count, const std::string& snapshot_id) { - service_->CleanupOldSnapshot(keep_count, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->CleanupOldSnapshot(keep_count, + snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->CleanupOldSnapshot(keep_count, snapshot_id); + } } tl::expected CallUploadSnapshotPayloadFile( const std::vector& data, const std::string& path, const std::string& local_filename, const std::string& snapshot_id) { - return service_->UploadSnapshotPayloadFile(data, path, local_filename, - snapshot_id); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->UploadSnapshotPayloadFile( + data, path, local_filename, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->UploadSnapshotPayloadFile( + data, path, local_filename, snapshot_id); + } } SnapshotObjectStore* GetSnapshotObjectStore() { @@ -149,12 +179,22 @@ class SnapshotChildProcessTest : public ::testing::Test { tl::expected CallPersistState( const std::string& snapshot_id) { - return service_->PersistState(snapshot_id); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->PersistState(snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->PersistState(snapshot_id); + } } tl::expected CallPersistState( const ha::SnapshotDescriptor& descriptor) { - return service_->PersistState(descriptor); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->PersistState(descriptor); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->PersistState(descriptor); + } } bool GetUseSnapshotBackupDir() { @@ -206,6 +246,47 @@ class SnapshotChildProcessTest : public ::testing::Test { return key + "_group"; } + private: + // Helper to create a temporary snapshot manager for tests + std::unique_ptr CreateTempSnapshotManager() { + EnsureSnapshotStores(); + + MasterSnapshotManagerOptions options; + options.enable_snapshot = true; + options.snapshot_interval_seconds = + service_->snapshot_interval_seconds_; + options.snapshot_child_timeout_seconds = + service_->snapshot_child_timeout_seconds_; + options.snapshot_retention_count = service_->snapshot_retention_count_; + options.snapshot_backup_dir = service_->snapshot_backup_dir_; + options.use_snapshot_backup_dir = service_->use_snapshot_backup_dir_; + options.snapshot_catalog_store_type = + service_->snapshot_catalog_store_type_; + options.snapshot_catalog_store_connstring = + service_->snapshot_catalog_store_connstring_; + options.ha_backend_type = service_->ha_backend_type_; + options.ha_backend_connstring = service_->ha_backend_connstring_; + options.cluster_id = service_->cluster_id_; + options.enable_ha = service_->enable_ha_; + + return std::make_unique( + service_.get(), options, service_->snapshot_mutex_, + service_->snapshot_object_store_.get(), + service_->snapshot_catalog_store_.get()); + } + + void EnsureSnapshotStores() { + if (!service_->snapshot_object_store_) { + service_->snapshot_object_store_ = SnapshotObjectStore::Create( + SnapshotObjectStoreType::LOCAL_FILE); + } + if (!service_->snapshot_catalog_store_ && + service_->snapshot_object_store_) { + service_->snapshot_catalog_store_ = + service_->CreateSnapshotCatalogStore(); + } + } + private: std::string tmp_dir_; }; From d409ed8986c2ada674c7fdd9a3f92bde018032d5 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Fri, 10 Jul 2026 18:37:43 +0800 Subject: [PATCH 47/57] [TENT] Add IntentType enum to Request for Transfer Intent API (#2810) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] Add IntentType enum to Request for Transfer Intent API Define standard intent categories (FOREGROUND_GET, BACKGROUND_PREFETCH, MIGRATION, CHECKPOINT, WEIGHT_LOADING, STAGING_INTERNAL) so TENT can identify a request's business semantics before scheduling. Changes: - types.h: add IntentType enum class + Request::intent_type field (default INTENT_UNSPEC, behavior byte-identical to today) - pybind.cpp: export IntentType enum, add intent_type/policy_name/ deadline_ns to Request constructor and as readwrite attributes - intent_type_test.cpp: 6 gtest cases covering defaults, assignment, integer values, field independence, copy, and batch usage Relates to: TENT roadmap "Transfer Intent API" * ci: retrigger CI * fix: keep intent type binding scoped * test: wire intent type coverage into cmake --------- Co-authored-by: 彦纾 Co-authored-by: Yanshu <237344440@qq.com> --- .../tent/include/tent/common/types.h | 11 +++ .../tent/src/python/pybind.cpp | 19 +++- .../tent/tests/CMakeLists.txt | 7 ++ .../tent/tests/intent_type_test.cpp | 95 +++++++++++++++++++ 4 files changed, 129 insertions(+), 3 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/intent_type_test.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index 4b7fa80503..a5b809bd7c 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -66,6 +66,16 @@ inline TransportType c_to_transport_hint(int v) { return static_cast(v); } +enum class IntentType : int { + INTENT_UNSPEC = 0, + FOREGROUND_GET, + BACKGROUND_PREFETCH, + MIGRATION, + CHECKPOINT, + WEIGHT_LOADING, + STAGING_INTERNAL, +}; + struct Request { enum OpCode { READ, WRITE }; OpCode opcode; @@ -86,6 +96,7 @@ struct Request { // (MLU = actual transfer time / available window) on completion; it does // not yet drive any admission or scheduling decision. See RFC #2519. uint64_t deadline_ns = 0; + IntentType intent_type = IntentType::INTENT_UNSPEC; }; enum TransferStatusEnum { diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index ac7c404b67..c1f2abe390 100644 --- a/mooncake-transfer-engine/tent/src/python/pybind.cpp +++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp @@ -302,6 +302,16 @@ PYBIND11_MODULE(tent, m) { .value("SUNRISE_LINK", TransportType::SUNRISE_LINK) .export_values(); + py::enum_(m, "IntentType") + .value("INTENT_UNSPEC", IntentType::INTENT_UNSPEC) + .value("FOREGROUND_GET", IntentType::FOREGROUND_GET) + .value("BACKGROUND_PREFETCH", IntentType::BACKGROUND_PREFETCH) + .value("MIGRATION", IntentType::MIGRATION) + .value("CHECKPOINT", IntentType::CHECKPOINT) + .value("WEIGHT_LOADING", IntentType::WEIGHT_LOADING) + .value("STAGING_INTERNAL", IntentType::STAGING_INTERNAL) + .export_values(); + py::enum_(m, "SegmentInfoType") .value("Memory", SegmentInfo::Type::Memory) .value("File", SegmentInfo::Type::File) @@ -324,7 +334,7 @@ PYBIND11_MODULE(tent, m) { size_t length, int priority, TransportType transport_hint, std::optional policy_name, - uint64_t deadline_ns) { + uint64_t deadline_ns, IntentType intent_type) { Request r; r.opcode = opcode; r.source = U64ToPtr(source); @@ -335,13 +345,15 @@ PYBIND11_MODULE(tent, m) { r.transport_hint = transport_hint; r.policy_name = std::move(policy_name); r.deadline_ns = deadline_ns; + r.intent_type = intent_type; return r; }), py::arg("opcode"), py::arg("source"), py::arg("target_id"), py::arg("target_offset"), py::arg("length"), py::arg("priority") = PRIO_HIGH, py::arg("transport_hint") = TransportType::UNSPEC, - py::arg("policy_name") = std::nullopt, py::arg("deadline_ns") = 0) + py::arg("policy_name") = std::nullopt, py::arg("deadline_ns") = 0, + py::arg("intent_type") = IntentType::INTENT_UNSPEC) .def_property( "opcode", [](const Request& r) { return r.opcode; }, [](Request& r, Request::OpCode op) { r.opcode = op; }) @@ -354,7 +366,8 @@ PYBIND11_MODULE(tent, m) { .def_readwrite("priority", &Request::priority) .def_readwrite("transport_hint", &Request::transport_hint) .def_readwrite("policy_name", &Request::policy_name) - .def_readwrite("deadline_ns", &Request::deadline_ns); + .def_readwrite("deadline_ns", &Request::deadline_ns) + .def_readwrite("intent_type", &Request::intent_type); py::class_(m, "TransferStatus") .def(py::init<>()) diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index c117b316b1..c8fcee00b5 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -189,6 +189,13 @@ target_include_directories(tent_transport_hint_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_transport_hint_test COMMAND tent_transport_hint_test) +add_executable(tent_intent_type_test intent_type_test.cpp) +target_link_libraries(tent_intent_type_test PRIVATE gtest gtest_main + tent_common) +target_include_directories(tent_intent_type_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_intent_type_test COMMAND tent_intent_type_test) + # ProgressWorker skeleton test: covers default-off behavior, event-driven # progress without poll-failover, and freeBatch races (issue #2116). add_executable(tent_progress_worker_test progress_worker_test.cpp) diff --git a/mooncake-transfer-engine/tent/tests/intent_type_test.cpp b/mooncake-transfer-engine/tent/tests/intent_type_test.cpp new file mode 100644 index 0000000000..ac56a2852e --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/intent_type_test.cpp @@ -0,0 +1,95 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Unit tests for IntentType enum and its integration with Request. + +#include + +#include + +#include "tent/common/types.h" + +namespace mooncake { +namespace tent { +namespace { + +TEST(IntentTypeTest, DefaultIsUnspec) { + Request r{}; + EXPECT_EQ(r.intent_type, IntentType::INTENT_UNSPEC); +} + +TEST(IntentTypeTest, AllValuesAssignable) { + Request r{}; + r.intent_type = IntentType::FOREGROUND_GET; + EXPECT_EQ(r.intent_type, IntentType::FOREGROUND_GET); + r.intent_type = IntentType::BACKGROUND_PREFETCH; + EXPECT_EQ(r.intent_type, IntentType::BACKGROUND_PREFETCH); + r.intent_type = IntentType::MIGRATION; + EXPECT_EQ(r.intent_type, IntentType::MIGRATION); + r.intent_type = IntentType::CHECKPOINT; + EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); + r.intent_type = IntentType::WEIGHT_LOADING; + EXPECT_EQ(r.intent_type, IntentType::WEIGHT_LOADING); + r.intent_type = IntentType::STAGING_INTERNAL; + EXPECT_EQ(r.intent_type, IntentType::STAGING_INTERNAL); +} + +TEST(IntentTypeTest, IntegerValues) { + EXPECT_EQ(static_cast(IntentType::INTENT_UNSPEC), 0); + EXPECT_EQ(static_cast(IntentType::FOREGROUND_GET), 1); + EXPECT_EQ(static_cast(IntentType::BACKGROUND_PREFETCH), 2); + EXPECT_EQ(static_cast(IntentType::MIGRATION), 3); + EXPECT_EQ(static_cast(IntentType::CHECKPOINT), 4); + EXPECT_EQ(static_cast(IntentType::WEIGHT_LOADING), 5); + EXPECT_EQ(static_cast(IntentType::STAGING_INTERNAL), 6); +} + +TEST(IntentTypeTest, DoesNotAffectOtherFields) { + Request r{}; + r.opcode = Request::READ; + r.priority = PRIO_LOW; + r.deadline_ns = 12345; + r.transport_hint = RDMA; + r.intent_type = IntentType::CHECKPOINT; + + EXPECT_EQ(r.opcode, Request::READ); + EXPECT_EQ(r.priority, PRIO_LOW); + EXPECT_EQ(r.deadline_ns, 12345u); + EXPECT_EQ(r.transport_hint, RDMA); + EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); +} + +TEST(IntentTypeTest, CopyPreservesIntentType) { + Request r{}; + r.intent_type = IntentType::WEIGHT_LOADING; + Request copy = r; + EXPECT_EQ(copy.intent_type, IntentType::WEIGHT_LOADING); +} + +TEST(IntentTypeTest, VectorOfRequests) { + std::vector batch(4); + batch[0].intent_type = IntentType::FOREGROUND_GET; + batch[1].intent_type = IntentType::BACKGROUND_PREFETCH; + batch[2].intent_type = IntentType::MIGRATION; + batch[3].intent_type = IntentType::INTENT_UNSPEC; + + EXPECT_EQ(batch[0].intent_type, IntentType::FOREGROUND_GET); + EXPECT_EQ(batch[1].intent_type, IntentType::BACKGROUND_PREFETCH); + EXPECT_EQ(batch[2].intent_type, IntentType::MIGRATION); + EXPECT_EQ(batch[3].intent_type, IntentType::INTENT_UNSPEC); +} + +} // namespace +} // namespace tent +} // namespace mooncake From 4188e3ae9923b93a74bf82d67be6a56d2ffb20e9 Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Fri, 10 Jul 2026 19:43:23 +0800 Subject: [PATCH 48/57] [TENT] admission queue: deadline proximity promotion for dispatch (#2814) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] admission queue: add deadline proximity promotion When a queued owner's remaining slack (deadline_ns - now) falls below the configurable promotion_slack_ns threshold, it is promoted to the front of the dispatch queue via stable_partition, ahead of owners with comfortable slack or no deadline. This complements the existing step-2 EDF ordering by dynamically boosting urgency as deadlines approach, ensuring near-deadline transfers get dispatched preferentially even if they were admitted after requests with later deadlines. Key design choices: - Opt-in: promotion_slack_ns defaults to 0 (disabled). - Requires deadline_aware = true (like all deadline features). - stable_partition preserves EDF order within promoted/non-promoted groups. - Composes with step-3 drop: promotion happens first, then infeasible owners are still dropped from the (now-reordered) front. - NowProvider reused from step-3 (falls back to steady_clock). Ref: RFC #2519 step 4 * fix: wire deadline promotion runtime config * fix: avoid deque stable partition in dispatch * perf: reuse deadline promotion scratch buffers * bench: add deadline promotion hot-path benchmark * perf: retain faster local partition buffers * perf: retain benchmarked stable partition --------- Co-authored-by: 彦纾 Co-authored-by: Yanshu <237344440@qq.com> --- .../include/tent/runtime/admission_queue.h | 8 + .../tent/src/runtime/admission_queue.cpp | 25 ++- .../tent/src/runtime/transfer_engine_impl.cpp | 6 + .../tent/tests/CMakeLists.txt | 8 + .../tent/tests/admission_queue_test.cpp | 191 ++++++++++++++++++ .../tent/tests/deadline_promotion_bench.cpp | 105 ++++++++++ 6 files changed, 340 insertions(+), 3 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index 8746bcd464..d74cb087e4 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -58,6 +58,14 @@ struct QueueLimits { // dropped instead of dispatched, and on_local_decode_suggested is raised so // the caller can recompute locally. Requires deadline_aware = true. double mlu_local_threshold{0.0}; + // Opt-in deadline proximity promotion. When > 0, pickForDispatch promotes + // queued owners whose remaining slack (deadline_ns - now) is below this + // threshold to the front of the dispatch queue, ahead of owners with more + // slack or no deadline. This dynamically boosts urgency as a deadline + // approaches, regardless of original admission order. Requires a + // NowProvider (via setDegradationPolicy) or defaults to steady_clock. + // 0 (default) disables promotion entirely. + uint64_t promotion_slack_ns{0}; }; struct QueueOwnerInput { diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index e13fdb0bb7..6eae246f95 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -242,9 +242,12 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( const bool drop_enabled = limits_.deadline_aware && limits_.mlu_local_threshold > 0.0 && static_cast(bandwidth_provider_); + const bool promotion_enabled = + limits_.deadline_aware && limits_.promotion_slack_ns > 0; + const bool need_now = drop_enabled || promotion_enabled; const double bw_bps = drop_enabled ? bandwidth_provider_() : 0.0; const uint64_t now_ns = - drop_enabled + need_now ? (now_provider_ ? now_provider_() : static_cast( @@ -254,13 +257,29 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( .count())) : 0; + // Deadline proximity promotion: partition fifo_ so owners with critical + // slack (deadline approaching within promotion_slack_ns) appear before + // owners with comfortable slack or no deadline. stable_partition preserves + // relative EDF order within each group. + if (promotion_enabled) { + std::stable_partition(fifo_.begin(), fifo_.end(), [&](QueueOwnerId id) { + auto it = owners_.find(id); + if (it == owners_.end() || it->second.state != QueueState::Queued) { + return false; + } + const uint64_t dl = it->second.request.deadline_ns; + if (dl == 0 || dl <= now_ns) return false; + return (dl - now_ns) < limits_.promotion_slack_ns; + }); + } + // Predicted MLU = predicted_transfer_time / remaining_window. Returns true // if the owner is predicted to miss its deadline hard enough to drop. auto shouldDrop = [&](const QueueOwner& owner) -> bool { if (!drop_enabled || bw_bps <= 0.0) return false; const uint64_t deadline_ns = owner.request.deadline_ns; - if (deadline_ns == 0) return false; // no deadline → never dropped - if (deadline_ns <= now_ns) return true; // already past → infeasible + if (deadline_ns == 0) return false; // no deadline + if (deadline_ns <= now_ns) return true; // already past const double window_s = (deadline_ns - now_ns) / 1e9; const double predicted_time_s = owner.request.length / bw_bps; const double mlu = predicted_time_s / window_s; diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 3629b8e5b5..6bee9e27aa 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -301,6 +301,12 @@ Status TransferEngineImpl::construct() { conf_->get("runtime_queue/staging_owner_reserve", 0UL); runtime_queue_config_.limits.staging_byte_reserve = conf_->get("runtime_queue/staging_byte_reserve", 0UL); + runtime_queue_config_.limits.deadline_aware = + conf_->get("runtime_queue/deadline_aware", false); + runtime_queue_config_.limits.mlu_local_threshold = + conf_->get("runtime_queue/mlu_local_threshold", 0.0); + runtime_queue_config_.limits.promotion_slack_ns = + conf_->get("runtime_queue/promotion_slack_ns", 0UL); runtime_queue_config_.max_dispatch_owners = conf_->get("runtime_queue/max_dispatch_owners", 64UL); runtime_queue_config_.max_dispatch_bytes = diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index c8fcee00b5..d04805f1ae 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -27,6 +27,14 @@ target_include_directories(admission_queue_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME admission_queue_test COMMAND admission_queue_test) +# Reproducible hot-path microbenchmark; intentionally not registered with +# ctest. Run manually when changing deadline promotion partitioning. +add_executable(deadline_promotion_bench deadline_promotion_bench.cpp + ../src/runtime/admission_queue.cpp) +target_link_libraries(deadline_promotion_bench PRIVATE tent_common) +target_include_directories(deadline_promotion_bench + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_executable(promotion_policy_test promotion_policy_test.cpp) target_link_libraries(promotion_policy_test PRIVATE tent_common gtest gtest_main) target_include_directories(promotion_policy_test diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index 083f36994c..30620a40c0 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -14,6 +14,7 @@ #include "tent/runtime/admission_queue.h" +#include #include #include @@ -595,6 +596,196 @@ TEST(AdmissionQueueTest, Step3NoDropWithoutBandwidthProvider) { EXPECT_TRUE(dropped.empty()); } +// --- Deadline proximity promotion (step 4) -------------------------------- + +QueueLimits promotionLimits(uint64_t slack_ns) { + QueueLimits limits{8, 1 << 20, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = slack_ns; + return limits; +} + +TEST(AdmissionQueueTest, PromotionDisabledKeepsEdfOrder) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2000), + makeOwnerWithDeadline(1, 16, 1500), + makeOwnerWithDeadline(2, 16, 1800)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 4096); + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionMovesUrgentOwnersToFront) { + LocalTransferAdmissionQueue queue(promotionLimits(500)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2000), + makeOwnerWithDeadline(1, 16, 1400), + makeOwnerWithDeadline(2, 16, 1300)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{3, 2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionReordersAcrossSeparateAdmits) { + LocalTransferAdmissionQueue queue(promotionLimits(2000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{5000}; }); + + std::vector ids; + auto s1 = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 10000)}), ids); + ASSERT_EQ(s1.code(), Status::Code::kOk); + auto s2 = queue.tryAdmit( + makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 6500)}), ids); + ASSERT_EQ(s2.code(), Status::Code::kOk); + auto s3 = queue.tryAdmit( + makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 6000)}), ids); + ASSERT_EQ(s3.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{3, 2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionSkipsNoDeadlineOwners) { + LocalTransferAdmissionQueue queue(promotionLimits(5000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 2, + {makeOwnerWithDeadline(0, 16, 0), + makeOwnerWithDeadline(1, 16, 2000)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionPreservesEdfWithinPromotedGroup) { + LocalTransferAdmissionQueue queue(promotionLimits(2000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2500), + makeOwnerWithDeadline(1, 16, 2200), + makeOwnerWithDeadline(2, 16, 2800)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{2, 1, 3}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionCoexistsWithStep3Drop) { + QueueLimits limits = promotionLimits(500); + limits.mlu_local_threshold = 1.5; + LocalTransferAdmissionQueue queue(limits); + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 1010), + makeOwnerWithDeadline(1, 16, 1400), + makeOwnerWithDeadline(2, 16, 5000)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + + const std::vector exp_pick{2, 3}; + const std::vector exp_drop{1}; + EXPECT_EQ(picked, exp_pick); + EXPECT_EQ(dropped, exp_drop); + EXPECT_EQ(hook_calls, 1); +} + +TEST(AdmissionQueueTest, PromotionWithAdvancingTime) { + QueueLimits limits = promotionLimits(500); + LocalTransferAdmissionQueue queue(limits); + + uint64_t fake_now = 1000; + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [&] { return fake_now; }); + + std::vector ids; + auto s1 = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1800)}), ids); + ASSERT_EQ(s1.code(), Status::Code::kOk); + auto s2 = queue.tryAdmit( + makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 1400)}), ids); + ASSERT_EQ(s2.code(), Status::Code::kOk); + auto s3 = queue.tryAdmit( + makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 3000)}), ids); + ASSERT_EQ(s3.code(), Status::Code::kOk); + + auto picked1 = queue.pickForDispatch(1, 1 << 20); + ASSERT_EQ(picked1.size(), 1u); + EXPECT_EQ(picked1[0], 2u); + + auto cstatus = queue.complete(2, TransferStatusEnum::COMPLETED); + ASSERT_EQ(cstatus.code(), Status::Code::kOk); + + fake_now = 1500; + auto picked2 = queue.pickForDispatch(2, 1 << 20); + const std::vector expected2{1, 3}; + EXPECT_EQ(picked2, expected2); +} + +TEST(AdmissionQueueTest, PromotionDisabledWithoutDeadlineAware) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = false; + limits.promotion_slack_ns = 5000; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 1200), + makeOwnerWithDeadline(1, 16, 5000), + makeOwnerWithDeadline(2, 16, 1100)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 4096); + const std::vector expected{1, 2, 3}; + EXPECT_EQ(picked, expected); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp b/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp new file mode 100644 index 0000000000..52def2ed0f --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp @@ -0,0 +1,105 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/admission_queue.h" + +namespace mooncake { +namespace tent { +namespace { + +constexpr uint64_t kNowNs = 1'000'000; + +QueueOwnerInput makeOwner(size_t task_id, bool urgent) { + QueueOwnerInput owner; + owner.owner_task_id = task_id; + owner.request.opcode = Request::READ; + owner.request.target_id = 1; + owner.request.length = 4096; + owner.request.deadline_ns = urgent ? kNowNs + 10'000 : kNowNs + 1'000'000; + return owner; +} + +double percentile(std::vector samples, double ratio) { + std::sort(samples.begin(), samples.end()); + const size_t index = + static_cast(ratio * static_cast(samples.size() - 1)); + return samples[index]; +} + +void runDepth(size_t depth) { + const size_t repeats = std::max(100, 200'000 / depth); + std::vector samples; + samples.reserve(repeats * depth); + + for (size_t repeat = 0; repeat < repeats; ++repeat) { + QueueLimits limits{depth, depth * 4096, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = 20'000; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return kNowNs; }); + + QueueSubmit submit; + submit.batch_token = repeat + 1; + submit.batch_slots_left = depth; + submit.owners.reserve(depth); + for (size_t i = 0; i < depth; ++i) { + // Interleave urgent and comfortable requests so every dispatch + // round exercises the stable partition. + submit.owners.push_back(makeOwner(i, i % 4 == 3)); + } + std::vector admitted; + if (!queue.tryAdmit(submit, admitted).ok()) std::abort(); + + for (size_t i = 0; i < depth; ++i) { + const auto start = std::chrono::steady_clock::now(); + auto picked = queue.pickForDispatch(1, 4096); + const auto end = std::chrono::steady_clock::now(); + if (picked.size() != 1 || + !queue.complete(picked[0], COMPLETED).ok()) { + std::abort(); + } + samples.push_back( + std::chrono::duration(end - start).count()); + } + if (!queue.retireBatch(submit.batch_token).ok()) std::abort(); + } + + const double mean = + std::accumulate(samples.begin(), samples.end(), 0.0) / samples.size(); + std::cout << depth << ',' << samples.size() << ',' << std::fixed + << std::setprecision(3) << mean << ',' + << percentile(samples, 0.50) << ',' << percentile(samples, 0.95) + << '\n'; +} + +} // namespace +} // namespace tent +} // namespace mooncake + +int main() { + std::cout << "queue_depth,samples,mean_us,p50_us,p95_us\n"; + for (size_t depth : {32, 64, 128, 256, 512}) { + mooncake::tent::runDepth(depth); + } + return 0; +} From 6894c2d77c9dfb7962d23b2a96c74cf635d80b5b Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Sat, 11 Jul 2026 11:32:18 +0800 Subject: [PATCH 49/57] [TransferEngine] Add show-link diagnostic tool for NIC topology (#2820) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [transfer-engine] Add show-link diagnostic tool for NIC topology inspection Add a connectivity diagnostic utility that displays local RDMA NIC information and topology selection matrix. This helps operators understand which NICs are discovered, their NUMA affinity, link speed, and how the topology-aware device selection maps storage types to NICs. Components: - show_links.h/cpp: Core logic (collectLocalNics, buildShowLinksReadable/Json) - show_link.cpp: CLI binary with --json/--discover_only flags - C API: showLinks() exposed via transfer_engine_c.h - rdma_context: Add activeWidth() accessor for bandwidth calculation - show_links_test.cpp: Basic gtest coverage Tested on 4-node H20 cluster (5 NICs per node, mlx5_0 + mlx5_bond_0..3): - NIC discovery: 5/5 devices found on both nodes - NUMA topology: Correctly distinguishes preferred vs available NICs - Cross-node probe: Verified at 0.5ms latency (3 consecutive runs) - Speed calculation: 200Gbps (HDR 50Gbps x 4 lanes) matches ibstat * style: clang-format show_links.cpp and transfer_engine.cpp * fix: respect show-link json output * fix: handle show-links before engine init * fix: make show-link discover-only discover topology * fix: show topology nics without rdma transport * fix: harden show-link C API and speed reporting * style: format show-links JSON assertion * fix: include integer types in show-links header --------- Co-authored-by: 彦纾 Co-authored-by: Yanshu <237344440@qq.com> --- .../example/CMakeLists.txt | 4 + .../example/show_link.cpp | 69 ++++++ mooncake-transfer-engine/include/show_links.h | 42 ++++ .../include/transfer_engine.h | 2 + .../include/transfer_engine_c.h | 3 + .../include/transfer_engine_impl.h | 1 + .../transport/rdma_transport/rdma_context.h | 2 + .../transport/rdma_transport/rdma_transport.h | 4 + mooncake-transfer-engine/src/show_links.cpp | 197 ++++++++++++++++++ .../src/transfer_engine.cpp | 16 ++ .../src/transfer_engine_c.cpp | 11 + .../transport/rdma_transport/rdma_context.cpp | 1 + mooncake-transfer-engine/tests/CMakeLists.txt | 4 + .../tests/show_links_test.cpp | 120 +++++++++++ 14 files changed, 476 insertions(+) create mode 100644 mooncake-transfer-engine/example/show_link.cpp create mode 100644 mooncake-transfer-engine/include/show_links.h create mode 100644 mooncake-transfer-engine/src/show_links.cpp create mode 100644 mooncake-transfer-engine/tests/show_links_test.cpp diff --git a/mooncake-transfer-engine/example/CMakeLists.txt b/mooncake-transfer-engine/example/CMakeLists.txt index 4adfcdba93..4aed903ea3 100644 --- a/mooncake-transfer-engine/example/CMakeLists.txt +++ b/mooncake-transfer-engine/example/CMakeLists.txt @@ -106,3 +106,7 @@ if(USE_CUDA AND BUILD_DEVICE_TRANSPORT_EXAMPLE) CUDA_ARCHITECTURES "80;90") endif() endif() + +add_executable(show_link ${WORKSPACE}/show_link.cpp) +target_link_libraries(show_link PUBLIC transfer_engine gflags::gflags + glog::glog) diff --git a/mooncake-transfer-engine/example/show_link.cpp b/mooncake-transfer-engine/example/show_link.cpp new file mode 100644 index 0000000000..ee0e3f0867 --- /dev/null +++ b/mooncake-transfer-engine/example/show_link.cpp @@ -0,0 +1,69 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include +#include + +#include "show_links.h" +#include "transfer_engine.h" + +DEFINE_string(metadata_server, "etcd://127.0.0.1:2379", + "Metadata server connection string"); +DEFINE_string(local_server_name, "", "Local server name (ip:port)"); +DEFINE_string(ip_or_host_name, "", "IP or hostname for RPC"); +DEFINE_int32(rpc_port, 12345, "RPC port"); +DEFINE_bool(json, false, "Output in JSON format"); +DEFINE_bool(discover_only, false, + "Only discover local topology without connecting"); + +using namespace mooncake; + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + + if (FLAGS_discover_only) { + auto engine = std::make_unique(/*auto_discover=*/false); + auto topology = engine->getLocalTopology(); + if (topology) { + topology->discover({}); + } + std::cout << engine->showLinks(FLAGS_json) << std::endl; + return 0; + } + + if (FLAGS_local_server_name.empty()) { + LOG(ERROR) << "Must specify --local_server_name or --discover_only"; + return 1; + } + + auto engine = std::make_unique(/*auto_discover=*/true); + int ret = engine->init(FLAGS_metadata_server, FLAGS_local_server_name, + FLAGS_ip_or_host_name, FLAGS_rpc_port); + if (ret) { + LOG(ERROR) << "Failed to initialize TransferEngine"; + return 1; + } + + engine->installTransport("rdma", nullptr); + + std::cout << engine->showLinks(FLAGS_json) << std::endl; + + engine->freeEngine(); + return 0; +} diff --git a/mooncake-transfer-engine/include/show_links.h b/mooncake-transfer-engine/include/show_links.h new file mode 100644 index 0000000000..41726ddef6 --- /dev/null +++ b/mooncake-transfer-engine/include/show_links.h @@ -0,0 +1,42 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_SHOW_LINKS_H_ +#define MOONCAKE_SHOW_LINKS_H_ + +#include +#include +#include + +namespace mooncake { + +struct NicDiagInfo { + std::string device_name; + int numa_node; + std::string gid; + int gid_index; + uint8_t port; + int active_speed; + int active_width; +}; + +class TransferEngineImpl; + +std::string buildShowLinksJson(TransferEngineImpl* impl); + +std::string buildShowLinksReadable(TransferEngineImpl* impl); + +} // namespace mooncake + +#endif // MOONCAKE_SHOW_LINKS_H_ diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 9fa1d813de..2b29c6c544 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -198,6 +198,8 @@ class TransferEngine { std::shared_ptr getLocalTopology(); + std::string showLinks(bool json = false) const; + private: std::shared_ptr impl_; std::shared_ptr impl_tent_; diff --git a/mooncake-transfer-engine/include/transfer_engine_c.h b/mooncake-transfer-engine/include/transfer_engine_c.h index 82baf8ee43..5a26905c9e 100644 --- a/mooncake-transfer-engine/include/transfer_engine_c.h +++ b/mooncake-transfer-engine/include/transfer_engine_c.h @@ -165,6 +165,9 @@ int freeBatchID(transfer_engine_t engine, batch_id_t batch_id); int syncSegmentCache(transfer_engine_t engine); +int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, + int json); + #ifdef __cplusplus } #endif // __cplusplus diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index c5d919d232..ae7b80c9b6 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -342,6 +342,7 @@ class TransferEngineImpl { } Transport* getTransport(const std::string& proto) { + if (!multi_transports_) return nullptr; return multi_transports_->getTransport(proto); } diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h index 162b7fc99a..07b9e9170a 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h @@ -181,6 +181,7 @@ class RdmaContext { uint8_t numLagPorts() const { return num_lag_ports_; } int activeSpeed() const { return active_speed_; } + int activeWidth() const { return active_width_; } ibv_mtu activeMTU() const { return active_mtu_; } @@ -232,6 +233,7 @@ class RdmaContext { uint16_t lid_ = 0; int gid_index_ = -1; int active_speed_ = -1; + int active_width_ = 1; ibv_mtu active_mtu_; uint8_t num_lag_ports_ = 0; // 0/1 = not in LAG; ≥2 = LAG active ibv_gid gid_; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h index e7a171b8b1..334bba20e6 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -135,6 +135,10 @@ class RdmaTransport : public Transport { int &buffer_id, int &device_id, int retry_cnt = 0); + const std::vector> &getContextList() const { + return context_list_; + } + private: std::vector> context_list_; std::shared_ptr local_topology_; diff --git a/mooncake-transfer-engine/src/show_links.cpp b/mooncake-transfer-engine/src/show_links.cpp new file mode 100644 index 0000000000..b50b8bfccd --- /dev/null +++ b/mooncake-transfer-engine/src/show_links.cpp @@ -0,0 +1,197 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "show_links.h" + +#include + +#include +#include + +#include "transfer_engine_impl.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_transport.h" + +namespace mooncake { + +static std::vector collectLocalNics(TransferEngineImpl* impl) { + std::vector nics; + auto* transport = impl->getTransport("rdma"); + if (!transport) { + auto topo = impl->getLocalTopology(); + if (!topo) return nics; + + for (const auto& hca : topo->getHcaList()) { + NicDiagInfo info; + info.device_name = hca; + info.numa_node = -1; + info.gid_index = -1; + info.port = 0; + info.active_speed = 0; + info.active_width = 0; + nics.push_back(info); + } + return nics; + } + + auto* rdma = static_cast(transport); + for (auto& ctx : rdma->getContextList()) { + NicDiagInfo info; + info.device_name = ctx->deviceName(); + info.numa_node = ctx->socketId(); + info.gid = ctx->gid(); + info.gid_index = ctx->gidIndex(); + info.port = ctx->portNum(); + info.active_speed = ctx->activeSpeed(); + info.active_width = ctx->activeWidth(); + nics.push_back(info); + } + return nics; +} + +static bool hasTransportDetails(const NicDiagInfo& nic) { + return nic.gid_index >= 0; +} + +static int computeSpeedGbps(int active_speed, int active_width) { + if (active_width <= 0) return 0; + + int lane_gbps = 0; + switch (active_speed) { + case 1: + lane_gbps = 2; + break; + case 2: + lane_gbps = 5; + break; + case 4: + lane_gbps = 10; + break; + case 8: + lane_gbps = 10; + break; + case 16: + lane_gbps = 14; + break; + case 32: + lane_gbps = 25; + break; + case 64: + lane_gbps = 50; + break; + case 128: + lane_gbps = 100; + break; + case 256: + lane_gbps = 200; + break; + default: + lane_gbps = std::max(0, active_speed); + break; + } + int lanes = 1; + switch (active_width) { + case 1: + lanes = 1; + break; + case 2: + lanes = 4; + break; + case 4: + lanes = 8; + break; + case 8: + lanes = 12; + break; + case 16: + lanes = 2; + break; + default: + lanes = 1; + break; + } + return lane_gbps * lanes; +} + +std::string buildShowLinksJson(TransferEngineImpl* impl) { + Json::Value root; + + auto nics = collectLocalNics(impl); + Json::Value nic_arr(Json::arrayValue); + for (auto& nic : nics) { + Json::Value n; + n["device"] = nic.device_name; + n["numa"] = nic.numa_node; + n["gid"] = nic.gid; + n["gid_index"] = nic.gid_index; + n["port"] = nic.port; + n["speed_gbps"] = computeSpeedGbps(nic.active_speed, nic.active_width); + n["source"] = hasTransportDetails(nic) ? "rdma_transport" : "topology"; + nic_arr.append(n); + } + root["local_nics"] = nic_arr; + + auto topo = impl->getLocalTopology(); + if (topo) { + root["topology"] = topo->toJson(); + } + + Json::StreamWriterBuilder builder; + builder["indentation"] = " "; + return Json::writeString(builder, root); +} + +std::string buildShowLinksReadable(TransferEngineImpl* impl) { + std::ostringstream os; + + auto nics = collectLocalNics(impl); + os << "=== Local NICs ===\n"; + if (nics.empty()) { + os << " (no RDMA devices found)\n"; + } + for (auto& nic : nics) { + os << " " << nic.device_name; + if (!hasTransportDetails(nic)) { + os << " (topology only; RDMA transport not initialized)\n"; + continue; + } + os << " NUMA=" << nic.numa_node << " GID=" << nic.gid + << " (idx=" << nic.gid_index << ")" + << " Port=" << (int)nic.port << " " + << computeSpeedGbps(nic.active_speed, nic.active_width) << "Gbps\n"; + } + + auto topo = impl->getLocalTopology(); + if (topo && !topo->empty()) { + os << "\n=== Topology (NIC Selection) ===\n"; + auto matrix = topo->getMatrix(); + for (auto& [location, entry] : matrix) { + os << " " << location << " -> preferred: ["; + for (size_t i = 0; i < entry.preferred_hca.size(); i++) { + if (i > 0) os << ", "; + os << entry.preferred_hca[i]; + } + os << "] available: ["; + for (size_t i = 0; i < entry.avail_hca.size(); i++) { + if (i > 0) os << ", "; + os << entry.avail_hca[i]; + } + os << "]\n"; + } + } + + return os.str(); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 7bd2ced5d2..99d5fcb91d 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -14,6 +14,7 @@ #ifndef USE_TENT #include "transfer_engine.h" +#include "show_links.h" #include "transfer_engine_impl.h" #include @@ -221,6 +222,12 @@ std::shared_ptr TransferEngine::getLocalTopology() { return impl_->getLocalTopology(); } +std::string TransferEngine::showLinks(bool json) const { + if (!impl_) return "{}"; + return json ? buildShowLinksJson(impl_.get()) + : buildShowLinksReadable(impl_.get()); +} + } // namespace mooncake #else #include "transfer_engine.h" @@ -229,6 +236,7 @@ std::shared_ptr TransferEngine::getLocalTopology() { #include "tent/common/config.h" #include +#include "show_links.h" namespace mooncake { @@ -670,5 +678,13 @@ void* TransferEngine::getBaseAddr() { return impl_->getBaseAddr(); } +std::string TransferEngine::showLinks(bool json) const { + if (use_tent_ || !impl_) { + return json ? "{}" : "(TENT mode or not initialized)"; + } + return json ? buildShowLinksJson(impl_.get()) + : buildShowLinksReadable(impl_.get()); +} + } // namespace mooncake #endif diff --git a/mooncake-transfer-engine/src/transfer_engine_c.cpp b/mooncake-transfer-engine/src/transfer_engine_c.cpp index b643a15386..944502610a 100644 --- a/mooncake-transfer-engine/src/transfer_engine_c.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_c.cpp @@ -14,6 +14,7 @@ #include "transfer_engine_c.h" +#include #include #include @@ -250,3 +251,13 @@ int syncSegmentCache(transfer_engine_t engine) { TransferEngine *native = (TransferEngine *)engine; return native->syncSegmentCache(); } + +int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, + int json) { + if (!engine || !buf_out || buf_len == 0) return -1; + + TransferEngine *native = (TransferEngine *)engine; + auto result = native->showLinks(json != 0); + snprintf(buf_out, buf_len, "%s", result.c_str()); + return 0; +} diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp index b208b69aef..aa44f57f81 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -1136,6 +1136,7 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, lid_ = attr.lid; active_mtu_ = attr.active_mtu; active_speed_ = attr.active_speed; + active_width_ = attr.active_width; { std::lock_guard guard(gid_lock_); gid_index_ = gid_index; diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 83b35b353b..6ab204f820 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -292,3 +292,7 @@ if(ENABLE_MULTI_PROTOCOL) gtest_main) add_test(NAME mp_transport_test COMMAND mp_transport_test) endif() + +add_executable(show_links_test ${WORKSPACE}/show_links_test.cpp) +target_link_libraries(show_links_test PUBLIC transfer_engine gtest gtest_main) +add_test(NAME show_links_test COMMAND show_links_test) diff --git a/mooncake-transfer-engine/tests/show_links_test.cpp b/mooncake-transfer-engine/tests/show_links_test.cpp new file mode 100644 index 0000000000..27706ddd6e --- /dev/null +++ b/mooncake-transfer-engine/tests/show_links_test.cpp @@ -0,0 +1,120 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include +#include +#include + +#include "transfer_engine.h" +#include "transfer_engine_c.h" + +using namespace mooncake; + +TEST(ShowLinksTest, NoInitReturnsEmptyOrPlaceholder) { + auto engine = std::make_unique(false); + auto result = engine->showLinks(); + EXPECT_FALSE(result.empty()); + + auto json_result = engine->showLinks(true); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(json_result.data(), + json_result.data() + json_result.size(), &root, + &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); +} + +TEST(ShowLinksTest, CApiSupportsJsonAndRejectsInvalidArguments) { + auto engine = std::make_unique(false); + char output[4096] = {}; + auto handle = reinterpret_cast(engine.get()); + + ASSERT_EQ(::showLinks(handle, output, sizeof(output), 1), 0); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE( + reader->parse(output, output + std::strlen(output), &root, &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); + + EXPECT_NE(::showLinks(nullptr, output, sizeof(output), 0), 0); + EXPECT_NE(::showLinks(handle, nullptr, sizeof(output), 0), 0); + EXPECT_NE(::showLinks(handle, output, 0, 0), 0); +} + +TEST(ShowLinksTest, AutoDiscoverShowsNics) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(); + EXPECT_NE(result.find("Local NICs"), std::string::npos); +} + +TEST(ShowLinksTest, OutputContainsTopologySection) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(); + // If RDMA devices exist, should show topology + // If not, gracefully show empty + EXPECT_FALSE(result.empty()); +} + +TEST(ShowLinksTest, JsonOutputIsValid) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(true); + EXPECT_FALSE(result.empty()); + + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(result.data(), result.data() + result.size(), + &root, &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); +} + +TEST(ShowLinksTest, TopologyOnlyNicsAppearWithoutTransport) { + auto engine = std::make_unique(false); + auto topology = engine->getLocalTopology(); + ASSERT_NE(topology, nullptr); + ASSERT_EQ(topology->parse("{\"cpu:0\" : [[\"erdma_0\"],[\"erdma_1\"]]}"), + 0); + + auto readable = engine->showLinks(); + EXPECT_NE(readable.find("erdma_0"), std::string::npos); + EXPECT_NE(readable.find("erdma_1"), std::string::npos); + EXPECT_NE(readable.find("topology only"), std::string::npos); + + auto json_result = engine->showLinks(true); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(json_result.data(), + json_result.data() + json_result.size(), &root, + &errors)) + << errors; + + ASSERT_TRUE(root["local_nics"].isArray()); + ASSERT_EQ(root["local_nics"].size(), 2u); + for (const auto& nic : root["local_nics"]) { + EXPECT_EQ(nic["source"].asString(), "topology"); + } +} From 29b494629e8ae69330241dc27a127ec4c2539bda Mon Sep 17 00:00:00 2001 From: Yanshu WANG Date: Sun, 12 Jul 2026 20:52:25 +0800 Subject: [PATCH 50/57] [TransferEngine] Add graceful shutdown for SIGTERM/SIGINT (#2812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TENT] admission queue: add deadline proximity promotion When a queued owner's remaining slack (deadline_ns - now) falls below the configurable promotion_slack_ns threshold, it is promoted to the front of the dispatch queue via stable_partition, ahead of owners with comfortable slack or no deadline. This complements the existing step-2 EDF ordering by dynamically boosting urgency as deadlines approach, ensuring near-deadline transfers get dispatched preferentially even if they were admitted after requests with later deadlines. Key design choices: - Opt-in: promotion_slack_ns defaults to 0 (disabled). - Requires deadline_aware = true (like all deadline features). - stable_partition preserves EDF order within promoted/non-promoted groups. - Composes with step-3 drop: promotion happens first, then infeasible owners are still dropped from the (now-reordered) front. - NowProvider reused from step-3 (falls back to steady_clock). Ref: RFC #2519 step 4 * [transfer-engine] Add graceful shutdown for SIGTERM/SIGINT/SIGABRT When a Transfer Engine process is killed by SIGTERM/SIGINT/SIGABRT, RDMA resources (QPs, MRs) are left dangling. The peer side's QP still references the now-invalid MR, causing RDMA READ to silently return undefined data instead of an error. This adds an opt-in enableGracefulShutdown() API that: - Registers atexit() to call freeEngine() on all active engines - Installs signal handlers that call exit(128+signo) to trigger atexit - Ensures the existing deconstruct() path (QP destroy + MR dereg) runs The API is opt-in to avoid interfering with applications that manage their own signal handlers (e.g., gRPC, Python interpreters). * fix: make graceful shutdown signal path safe * fix: cover tent shutdown and preserve abort semantics * fix: preserve graceful shutdown across moves * fix: avoid post-fork shutdown handler hang --------- Co-authored-by: 彦纾 Co-authored-by: Yanshu <237344440@qq.com> --- .../include/graceful_shutdown.h | 41 ++++ .../include/transfer_engine.h | 9 +- .../include/transfer_engine_c.h | 1 + .../src/graceful_shutdown.cpp | 181 ++++++++++++++++++ .../src/transfer_engine.cpp | 157 +++++++++++++++ .../src/transfer_engine_c.cpp | 5 + mooncake-transfer-engine/tests/CMakeLists.txt | 5 + .../tests/graceful_shutdown_test.cpp | 156 +++++++++++++++ 8 files changed, 553 insertions(+), 2 deletions(-) create mode 100644 mooncake-transfer-engine/include/graceful_shutdown.h create mode 100644 mooncake-transfer-engine/src/graceful_shutdown.cpp create mode 100644 mooncake-transfer-engine/tests/graceful_shutdown_test.cpp diff --git a/mooncake-transfer-engine/include/graceful_shutdown.h b/mooncake-transfer-engine/include/graceful_shutdown.h new file mode 100644 index 0000000000..41a637c382 --- /dev/null +++ b/mooncake-transfer-engine/include/graceful_shutdown.h @@ -0,0 +1,41 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_GRACEFUL_SHUTDOWN_H_ +#define MOONCAKE_GRACEFUL_SHUTDOWN_H_ + +#include + +namespace mooncake { + +class TransferEngineImpl; + +class ShutdownToken { + public: + virtual ~ShutdownToken() = default; + + virtual void shutdown() = 0; + + virtual void detach() = 0; +}; + +void registerTokenForShutdown(std::shared_ptr token); + +void registerEngineForShutdown(std::shared_ptr impl); + +void installGracefulShutdownHandlers(); + +} // namespace mooncake + +#endif // MOONCAKE_GRACEFUL_SHUTDOWN_H_ diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 2b29c6c544..afa9fce2fb 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -15,12 +15,15 @@ #ifndef MULTI_TRANSFER_ENGINE_H_ #define MULTI_TRANSFER_ENGINE_H_ +#include + #include "memory_location.h" #include "multi_transport.h" #include "transfer_metadata.h" #include "transport/transport.h" namespace mooncake { +class ShutdownToken; class TransferEngineImpl; namespace tent { class TransferEngine; @@ -72,9 +75,9 @@ class TransferEngine { TransferEngine(bool auto_discover, const std::vector& filter); - TransferEngine(TransferEngine&&) = default; + TransferEngine(TransferEngine&& other) noexcept; - TransferEngine& operator=(TransferEngine&&) = default; + TransferEngine& operator=(TransferEngine&& other) noexcept; ~TransferEngine(); @@ -198,11 +201,13 @@ class TransferEngine { std::shared_ptr getLocalTopology(); + void enableGracefulShutdown(); std::string showLinks(bool json = false) const; private: std::shared_ptr impl_; std::shared_ptr impl_tent_; + std::shared_ptr shutdown_token_; bool use_tent_{false}; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transfer_engine_c.h b/mooncake-transfer-engine/include/transfer_engine_c.h index 5a26905c9e..23c7d7f703 100644 --- a/mooncake-transfer-engine/include/transfer_engine_c.h +++ b/mooncake-transfer-engine/include/transfer_engine_c.h @@ -165,6 +165,7 @@ int freeBatchID(transfer_engine_t engine, batch_id_t batch_id); int syncSegmentCache(transfer_engine_t engine); +void enableGracefulShutdown(transfer_engine_t engine); int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, int json); diff --git a/mooncake-transfer-engine/src/graceful_shutdown.cpp b/mooncake-transfer-engine/src/graceful_shutdown.cpp new file mode 100644 index 0000000000..8fbf2f0cf0 --- /dev/null +++ b/mooncake-transfer-engine/src/graceful_shutdown.cpp @@ -0,0 +1,181 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "graceful_shutdown.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "transfer_engine_impl.h" + +namespace mooncake { + +namespace { + +std::mutex g_registry_mutex; +std::vector> g_tokens; +std::atomic g_cleanup_started{false}; + +std::mutex g_install_mutex; +bool g_handlers_installed = false; +volatile sig_atomic_t g_handlers_pid = 0; +bool g_atexit_registered = false; +int g_signal_pipe[2] = {-1, -1}; +volatile sig_atomic_t g_signal_seen = 0; + +class TransferEngineImplShutdownToken : public ShutdownToken { + public: + explicit TransferEngineImplShutdownToken( + std::shared_ptr impl) + : impl_(std::move(impl)) {} + + void shutdown() override { + if (auto impl = impl_.lock()) { + impl->freeEngine(); + } + } + + void detach() override { impl_.reset(); } + + private: + std::weak_ptr impl_; +}; + +std::vector> collectTokensForCleanup() { + std::lock_guard lock(g_registry_mutex); + auto tokens = std::move(g_tokens); + g_tokens.clear(); + return tokens; +} + +void cleanupEngines() { + if (g_cleanup_started.exchange(true)) return; + + auto tokens = collectTokensForCleanup(); + for (auto& token : tokens) { + token->shutdown(); + } +} + +void atexitCleanup() { cleanupEngines(); } + +void signalWatcher() { + unsigned char signal_byte = 0; + ssize_t bytes_read = 0; + do { + bytes_read = read(g_signal_pipe[0], &signal_byte, sizeof(signal_byte)); + } while (bytes_read < 0 && errno == EINTR); + + if (bytes_read == static_cast(sizeof(signal_byte))) { + cleanupEngines(); + _Exit(128 + static_cast(signal_byte)); + } + + _Exit(1); +} + +bool startSignalWatcherLocked(pid_t current_pid) { + if (g_signal_pipe[0] >= 0) close(g_signal_pipe[0]); + if (g_signal_pipe[1] >= 0) close(g_signal_pipe[1]); + g_signal_pipe[0] = -1; + g_signal_pipe[1] = -1; + g_signal_seen = 0; + + if (pipe(g_signal_pipe) != 0) return false; + + try { + std::thread(signalWatcher).detach(); + } catch (...) { + close(g_signal_pipe[0]); + close(g_signal_pipe[1]); + g_signal_pipe[0] = -1; + g_signal_pipe[1] = -1; + return false; + } + + g_handlers_pid = static_cast(current_pid); + return true; +} + +void shutdownSignalHandler(int signo) { + // After fork(), only the calling thread survives. If the parent had already + // installed handlers, the child inherits the handler and pipe fds but not + // the watcher thread. Exit directly instead of blocking forever in pause(). + if (g_handlers_pid != static_cast(getpid())) { + _Exit(128 + signo); + } + + if (g_signal_seen == 0) { + g_signal_seen = signo; + unsigned char signal_byte = static_cast(signo); + if (g_signal_pipe[1] < 0 || + write(g_signal_pipe[1], &signal_byte, sizeof(signal_byte)) != + static_cast(sizeof(signal_byte))) { + _Exit(128 + signo); + } + } + + for (;;) { + pause(); + } +} + +} // namespace + +void registerTokenForShutdown(std::shared_ptr token) { + std::lock_guard lock(g_registry_mutex); + g_tokens.push_back(std::move(token)); +} + +void registerEngineForShutdown(std::shared_ptr impl) { + registerTokenForShutdown( + std::make_shared(std::move(impl))); +} + +void installGracefulShutdownHandlers() { + std::lock_guard lock(g_install_mutex); + pid_t current_pid = getpid(); + if (g_handlers_installed && + g_handlers_pid == static_cast(current_pid)) + return; + + if (!g_atexit_registered) { + atexit(atexitCleanup); + g_atexit_registered = true; + } + + if (!startSignalWatcherLocked(current_pid)) return; + + struct sigaction sa{}; + sa.sa_handler = shutdownSignalHandler; + sigemptyset(&sa.sa_mask); + sigaddset(&sa.sa_mask, SIGTERM); + sigaddset(&sa.sa_mask, SIGINT); + sa.sa_flags = 0; + sigaction(SIGTERM, &sa, nullptr); + sigaction(SIGINT, &sa, nullptr); + + g_handlers_installed = true; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 99d5fcb91d..a24a66dbb8 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -16,9 +16,52 @@ #include "transfer_engine.h" #include "show_links.h" #include "transfer_engine_impl.h" +#include "graceful_shutdown.h" +#include #include namespace mooncake { +namespace { + +class TransferEngineShutdownToken : public ShutdownToken { + public: + explicit TransferEngineShutdownToken(TransferEngine* engine) + : engine_(engine) {} + + void shutdown() override { + TransferEngine* engine = nullptr; + { + std::lock_guard lock(mutex_); + engine = engine_; + engine_ = nullptr; + } + if (engine) engine->freeEngine(); + } + + void detach() override { + std::lock_guard lock(mutex_); + engine_ = nullptr; + } + + private: + std::mutex mutex_; + TransferEngine* engine_; +}; + +std::shared_ptr registerTransferEngineShutdownToken( + TransferEngine* engine) { + auto token = std::make_shared(engine); + registerTokenForShutdown(token); + return token; +} + +void detachShutdownToken(std::shared_ptr& token) { + if (!token) return; + token->detach(); + token.reset(); +} + +} // namespace TransferEngine::TransferEngine(bool auto_discover) : impl_(std::make_shared(auto_discover)) {} @@ -27,6 +70,33 @@ TransferEngine::TransferEngine(bool auto_discover, const std::vector& filter) : impl_(std::make_shared(auto_discover, filter)) {} +TransferEngine::TransferEngine(TransferEngine&& other) noexcept + : impl_(std::move(other.impl_)), + impl_tent_(std::move(other.impl_tent_)), + use_tent_(other.use_tent_) { + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } +} + +TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { + if (this == &other) return *this; + freeEngine(); + impl_ = std::move(other.impl_); + impl_tent_ = std::move(other.impl_tent_); + use_tent_ = other.use_tent_; + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } + return *this; +} + TransferEngine::~TransferEngine() { freeEngine(); } int TransferEngine::init(const std::string& metadata_conn_string, @@ -38,6 +108,7 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (impl_) { impl_->freeEngine(); impl_.reset(); @@ -222,6 +293,13 @@ std::shared_ptr TransferEngine::getLocalTopology() { return impl_->getLocalTopology(); } +void TransferEngine::enableGracefulShutdown() { + if (!shutdown_token_) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + } + installGracefulShutdownHandlers(); +} + std::string TransferEngine::showLinks(bool json) const { if (!impl_) return "{}"; return json ? buildShowLinksJson(impl_.get()) @@ -235,10 +313,53 @@ std::string TransferEngine::showLinks(bool json) const { #include "tent/transfer_engine.h" #include "tent/common/config.h" +#include #include +#include "graceful_shutdown.h" #include "show_links.h" namespace mooncake { +namespace { + +class TransferEngineShutdownToken : public ShutdownToken { + public: + explicit TransferEngineShutdownToken(TransferEngine* engine) + : engine_(engine) {} + + void shutdown() override { + TransferEngine* engine = nullptr; + { + std::lock_guard lock(mutex_); + engine = engine_; + engine_ = nullptr; + } + if (engine) engine->freeEngine(); + } + + void detach() override { + std::lock_guard lock(mutex_); + engine_ = nullptr; + } + + private: + std::mutex mutex_; + TransferEngine* engine_; +}; + +std::shared_ptr registerTransferEngineShutdownToken( + TransferEngine* engine) { + auto token = std::make_shared(engine); + registerTokenForShutdown(token); + return token; +} + +void detachShutdownToken(std::shared_ptr& token) { + if (!token) return; + token->detach(); + token.reset(); +} + +} // namespace TransferEngine::TransferEngine(bool auto_discover) { if (getenv("MC_USE_TENT") || getenv("MC_USE_TEV1")) { @@ -259,6 +380,34 @@ TransferEngine::TransferEngine(bool auto_discover, } } +TransferEngine::TransferEngine(TransferEngine&& other) noexcept + : impl_(std::move(other.impl_)), + impl_tent_(std::move(other.impl_tent_)), + shutdown_token_(nullptr), + use_tent_(other.use_tent_) { + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } +} + +TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { + if (this == &other) return *this; + freeEngine(); + impl_ = std::move(other.impl_); + impl_tent_ = std::move(other.impl_tent_); + use_tent_ = other.use_tent_; + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } + return *this; +} + TransferEngine::~TransferEngine() { freeEngine(); } static std::pair parseConnectionStringInternal( @@ -308,6 +457,7 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (!use_tent_ && impl_) { impl_->freeEngine(); impl_.reset(); @@ -678,6 +828,13 @@ void* TransferEngine::getBaseAddr() { return impl_->getBaseAddr(); } +void TransferEngine::enableGracefulShutdown() { + if (!shutdown_token_) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + } + installGracefulShutdownHandlers(); +} + std::string TransferEngine::showLinks(bool json) const { if (use_tent_ || !impl_) { return json ? "{}" : "(TENT mode or not initialized)"; diff --git a/mooncake-transfer-engine/src/transfer_engine_c.cpp b/mooncake-transfer-engine/src/transfer_engine_c.cpp index 944502610a..5ed7e99d0a 100644 --- a/mooncake-transfer-engine/src/transfer_engine_c.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_c.cpp @@ -252,6 +252,11 @@ int syncSegmentCache(transfer_engine_t engine) { return native->syncSegmentCache(); } +void enableGracefulShutdown(transfer_engine_t engine) { + TransferEngine *native = (TransferEngine *)engine; + native->enableGracefulShutdown(); +} + int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, int json) { if (!engine || !buf_out || buf_len == 0) return -1; diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 6ab204f820..f453f73de3 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -293,6 +293,11 @@ if(ENABLE_MULTI_PROTOCOL) add_test(NAME mp_transport_test COMMAND mp_transport_test) endif() +add_executable(graceful_shutdown_test ${WORKSPACE}/graceful_shutdown_test.cpp) +target_link_libraries(graceful_shutdown_test PUBLIC transfer_engine gtest + gtest_main) +add_test(NAME graceful_shutdown_test COMMAND graceful_shutdown_test) + add_executable(show_links_test ${WORKSPACE}/show_links_test.cpp) target_link_libraries(show_links_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME show_links_test COMMAND show_links_test) diff --git a/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp new file mode 100644 index 0000000000..f0741c82bf --- /dev/null +++ b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp @@ -0,0 +1,156 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include + +#include "transfer_engine.h" + +using namespace mooncake; + +namespace { + +void waitChildWithTimeout(pid_t pid, int* status) { + for (int i = 0; i < 50; ++i) { + pid_t ret = waitpid(pid, status, WNOHANG); + ASSERT_NE(ret, -1) << "waitpid() failed"; + if (ret == pid) return; + usleep(100000); + } + kill(pid, SIGKILL); + waitpid(pid, status, 0); + FAIL() << "child did not exit before timeout"; +} + +} // namespace + +TEST(GracefulShutdownTest, SigtermTriggersCleanExit) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) + << "Child did not exit normally (signaled: " << WIFSIGNALED(status) + << ")"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, SigintTriggersCleanExit) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGINT); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) << "Child did not exit normally"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGINT); +} + +TEST(GracefulShutdownTest, IdempotentEnable) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + engine->enableGracefulShutdown(); + engine->enableGracefulShutdown(); +} + +TEST(GracefulShutdownTest, EngineDestroyedBeforeSignal) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + } + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, ForkAfterInstallDoesNotHangChildSignal) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) + << "Child did not exit normally (signaled: " << WIFSIGNALED(status) + << ")"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, MultipleEngines) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine1 = std::make_unique(false); + auto engine2 = std::make_unique(false); + engine1->enableGracefulShutdown(); + engine2->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} From 721ec57bcf0aa56d921ad95434af74c5b0b41924 Mon Sep 17 00:00:00 2001 From: "Guocheng(Eric) Song" Date: Sun, 12 Jul 2026 20:56:07 +0800 Subject: [PATCH 51/57] [TransferEngine] Validate batch memory registrations (#2854) --- .../src/transfer_engine_impl.cpp | 27 +++++++++++- .../tests/transport_uint_test.cpp | 43 ++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index d1514bf210..fe34f1841e 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -742,7 +742,32 @@ int TransferEngineImpl::mp_unregisterLocalMemory( int TransferEngineImpl::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { - for (auto& buffer : buffer_list) { + std::vector sorted_buffers = buffer_list; + std::sort(sorted_buffers.begin(), sorted_buffers.end(), + [](const BufferEntry& lhs, const BufferEntry& rhs) { + return reinterpret_cast(lhs.addr) < + reinterpret_cast(rhs.addr); + }); + + for (size_t i = 0; i < sorted_buffers.size(); ++i) { + const auto& buffer = sorted_buffers[i]; + if (buffer.length == 0) { + LOG(ERROR) + << "Transfer Engine does not support zero length memory region"; + return ERR_INVALID_ARGUMENT; + } + + if (i > 0) { + const auto& previous = sorted_buffers[i - 1]; + auto address = reinterpret_cast(buffer.addr); + auto previous_address = reinterpret_cast(previous.addr); + if (address - previous_address < previous.length) { + LOG(ERROR) << "Transfer Engine does not support overlapped " + "memory region"; + return ERR_ADDRESS_OVERLAPPED; + } + } + if (checkOverlap(buffer.addr, buffer.length)) { LOG(ERROR) << "Transfer Engine does not support overlapped memory region"; diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index 599617c6e6..6c1950fe98 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -171,9 +172,49 @@ TEST_F(TransportTest, ReadEmptyFile) { close(fd); } + +TEST_F(TransportTest, RegisterLocalMemoryBatchRejectsOverlappingBuffers) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data() + 64, 128}, + {buffer.data(), 128}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), + ERR_ADDRESS_OVERLAPPED); +} + +TEST_F(TransportTest, RegisterLocalMemoryBatchRejectsZeroLengthBuffer) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data(), 0}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), + ERR_INVALID_ARGUMENT); +} + +TEST_F(TransportTest, RegisterLocalMemoryBatchAllowsAdjacentBuffers) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data() + 128, 128}, + {buffer.data(), 128}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); +} } // namespace mooncake int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} From 2dfa138ca4b9375c17419c78f72963617f52559c Mon Sep 17 00:00:00 2001 From: Zoee <30841158+n-WN@users.noreply.github.com> Date: Sun, 12 Jul 2026 22:03:45 +0800 Subject: [PATCH 52/57] [TE] Make ThreadLocalStorage per-instance and reclaim per-thread holders (#2842) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [TE] Make ThreadLocalStorage per-instance and reclaim per-thread holders Fixes #2717. The previous implementation kept its thread-local slot as a 'thread_local static' member of the class template — one slot per template instantiation, shared by every ThreadLocalStorage instance in the process — so two instances (e.g. two engines' SegmentManager remote-desc caches) aliased each other's per-thread state. The slot was also a raw pointer to a heap holder that nothing ever deleted, so the deregistration logic in the holder destructor was unreachable and every thread leaked one holder per instantiation. Rework the storage around a per-thread map keyed by a process-monotonic instance id (never reused, so a recycled allocation cannot alias a stale entry), with a one-entry cache in front so the common get() stays one thread_local access plus a compare. Per-thread values are owned by the thread and destroyed at thread exit. A control block jointly owned by the storage and every thread node makes both teardown orders safe: a thread exiting first deregisters from the registry; an owner destroyed first marks the block dead and exiting threads skip the registry. forEach() runs under the registry mutex and visits exactly the live registered values. Validation (96-core box, GCC 13): - New thread_local_storage_test: 8/8. Against the previous implementation the same binary fails 5/7 (aliasing, deregistration, resurrection, forEach, churn) and LeakSanitizer reports 288 bytes in 18 allocations from get(). - ASAN/UBSAN and TSAN clean for 50 repeated runs each (both teardown orders exercised). - Full TENT suite 23/23 (SegmentManager's tl_remote_cache_ exercises the storage end-to-end). - get() hot path at -O3: 0.33-0.41 ns before, 0.66-0.69 ns after (+~0.3 ns; the caller's next step is a clock_gettime and map lookup). * review: null-check control in sweepDeadNodes Unreachable today (the sweep runs before try_emplace on the same thread and control is assigned immediately after emplace by a nothrow shared_ptr copy), but consistent with ~ThreadNode's defensive check and robust to future reordering. --- .../common/concurrent/thread_local_storage.h | 177 ++++++++++--- .../tent/tests/CMakeLists.txt | 6 + .../tent/tests/thread_local_storage_test.cpp | 234 ++++++++++++++++++ 3 files changed, 385 insertions(+), 32 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h b/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h index ab2fc36937..63b28accd1 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h +++ b/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h @@ -15,70 +15,183 @@ #ifndef TENT_THREAD_LOCAL_STORAGE_H #define TENT_THREAD_LOCAL_STORAGE_H +#include +#include #include +#include #include -#include +#include #include namespace mooncake { namespace tent { +namespace detail { +// Hoisted out of the class template so ids are unique across ALL +// ThreadLocalStorage instantiations, not just within one T. The per-thread +// maps are per-T, so per-T uniqueness would suffice today — global +// uniqueness removes the aliasing hazard if the per-thread state is ever +// shared across instantiations. +inline uint64_t nextThreadLocalStorageId() { + static std::atomic counter{1}; + return counter.fetch_add(1, std::memory_order_relaxed); +} +} // namespace detail + +// Per-instance thread-local storage (#2717). +// +// Each (ThreadLocalStorage instance, thread) pair owns a distinct T. The +// previous implementation kept one `thread_local` slot per template +// instantiation, so all instances of ThreadLocalStorage in a process +// aliased each other's per-thread state, and its heap-allocated holders were +// never destroyed (the deregistration path was unreachable). +// +// Lifetime rules: +// - A thread's T values are destroyed when the thread exits. +// - Destroying the storage does not destroy other threads' values; they are +// orphaned (at most one T per thread per destroyed storage) and reclaimed +// at those threads' exit. Instance ids are process-monotonic and never +// reused, so an orphaned value can never be served to a new instance that +// happens to reuse the same address. +// - Both teardown orders are safe: a thread exiting while the owner is alive +// deregisters its value from the owner's registry; an owner destroyed +// while threads are alive marks the jointly-owned control block dead, and +// those threads skip deregistration at exit. Orphans are additionally +// swept opportunistically: the next first-use get() of ANY instance on +// that thread reclaims all of the thread's dead-owner values, so orphan +// count stays bounded by live instances even under storage churn. +// +// Precondition: callers must keep the storage alive across every get() / +// forEach() call (the usual member-of-owner pattern satisfies this); the +// destructor may run concurrently only with other threads' exits, not with +// their accesses. +// +// Concurrency: +// - get() is lock-free after the first call per (instance, thread): one +// thread_local access plus an id compare on the hot path. +// - forEach() runs under the registry mutex and visits exactly the values of +// live registered threads. It synchronizes registry membership only — if +// owning threads mutate their T concurrently, the callback observes those +// fields with whatever synchronization T itself provides (unchanged from +// the previous implementation). template class ThreadLocalStorage { public: ThreadLocalStorage() = default; - ~ThreadLocalStorage() = default; - // Disable copy/move + ~ThreadLocalStorage() { + std::lock_guard lock(control_->mutex); + control_->owner_alive.store(false, std::memory_order_release); + control_->values.clear(); + } + ThreadLocalStorage(const ThreadLocalStorage&) = delete; ThreadLocalStorage& operator=(const ThreadLocalStorage&) = delete; - // Access the thread-local instance (lock-free) + // Access the calling thread's instance, constructing it on first use. T& get() { - if (!instance_) { - instance_ = new InstanceHolder(this); + ThreadState& state = threadState(); + if (state.cached_id == id_) return *state.cached_value; + auto it = state.nodes.find(id_); + if (it == state.nodes.end()) { + // First use of this instance on this thread — the cold path. + // Piggyback a sweep of values whose owners are gone, so a + // long-lived thread does not accumulate one orphan per + // destroyed storage it ever touched (their T contents would + // otherwise stay pinned until thread exit). + sweepDeadNodes(state); + it = state.nodes.try_emplace(id_).first; + ThreadNode& node = it->second; + node.control = control_; + // The caller holds a reference to *this, so the owner is alive + // and registration cannot race ~ThreadLocalStorage's clear(). + std::lock_guard lock(control_->mutex); + control_->values.insert(&node.value); } - return instance_->value; + state.cached_id = id_; + state.cached_value = &it->second.value; + return *state.cached_value; } - // Safe iteration over all instances (with locking) + // Safe iteration over the values of all live threads that have called + // get() on this instance. void forEach(const std::function& fn) { - std::lock_guard lock(global_mutex_); - for (auto* inst : instances_) { - fn(inst->value); + std::lock_guard lock(control_->mutex); + for (T* value : control_->values) { + fn(*value); } } private: - struct InstanceHolder { - T value; - ThreadLocalStorage* owner; + // Shared between the owner and every thread node so that whichever side + // is torn down last still has a valid registry (or a dead flag) to look + // at. + struct ControlBlock { + std::mutex mutex; + std::unordered_set values; + // Atomic so the orphan sweep can test liveness without taking the + // mutex of every node it scans. + std::atomic owner_alive{true}; + }; - InstanceHolder(ThreadLocalStorage* owner) : owner(owner) { - std::lock_guard lock(owner->global_mutex_); - owner->instances_.insert(this); - } + struct ThreadNode { + T value{}; + std::shared_ptr control; - ~InstanceHolder() { - std::lock_guard lock(owner->global_mutex_); - owner->instances_.erase(this); + ThreadNode() = default; + ThreadNode(const ThreadNode&) = delete; + ThreadNode& operator=(const ThreadNode&) = delete; + + ~ThreadNode() { + if (!control) return; + std::lock_guard lock(control->mutex); + if (control->owner_alive.load(std::memory_order_acquire)) + control->values.erase(&value); } }; - // Thread-local pointer to the per-thread instance - thread_local static InstanceHolder* instance_; + struct ThreadState { + // Node-based map: ThreadNode addresses are stable across rehash, + // which the registry and the one-entry cache below rely on. + std::unordered_map nodes; + // One-entry cache so the common get() is a single thread_local + // access plus a compare. Entries live until thread exit, so the + // cached pointer cannot dangle while the id matches. + uint64_t cached_id = 0; // ids start at 1; 0 never matches + T* cached_value = nullptr; + }; - // Global list of all thread instances - std::unordered_set instances_; - std::mutex global_mutex_; -}; + static void sweepDeadNodes(ThreadState& state) { + for (auto it = state.nodes.begin(); it != state.nodes.end();) { + // A null control cannot be observed today (the sweep runs before + // try_emplace on the same thread, and control is assigned + // immediately after emplace by a nothrow shared_ptr copy), but + // check it for consistency with ~ThreadNode and robustness to + // reordering. + if (!it->second.control || !it->second.control->owner_alive.load( + std::memory_order_acquire)) { + if (state.cached_value == &it->second.value) { + state.cached_id = 0; + state.cached_value = nullptr; + } + // ~ThreadNode sees the dead owner and skips the registry. + it = state.nodes.erase(it); + } else { + ++it; + } + } + } -// Definition of thread_local variable (must be outside the class) -template -thread_local typename ThreadLocalStorage::InstanceHolder* - ThreadLocalStorage::instance_ = nullptr; + static ThreadState& threadState() { + thread_local ThreadState state; + return state; + } + + const uint64_t id_ = detail::nextThreadLocalStorageId(); + std::shared_ptr control_ = std::make_shared(); +}; } // namespace tent } // namespace mooncake -#endif // TENT_THREAD_LOCAL_STORAGE_H \ No newline at end of file +#endif // TENT_THREAD_LOCAL_STORAGE_H diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index d04805f1ae..231e46ea0d 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -41,6 +41,12 @@ target_include_directories(promotion_policy_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME promotion_policy_test COMMAND promotion_policy_test) +add_executable(thread_local_storage_test thread_local_storage_test.cpp) +target_link_libraries(thread_local_storage_test PRIVATE gtest gtest_main) +target_include_directories(thread_local_storage_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME thread_local_storage_test COMMAND thread_local_storage_test) + add_executable(tent_ip_utils_test ip_utils_test.cpp) target_link_libraries(tent_ip_utils_test PRIVATE tent_common gtest gtest_main) target_include_directories(tent_ip_utils_test diff --git a/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp new file mode 100644 index 0000000000..cb6921a461 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp @@ -0,0 +1,234 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression tests for #2717: ThreadLocalStorage used one `thread_local` +// slot per template instantiation (all instances aliased each other's +// per-thread state) and never destroyed its per-thread holders (the +// deregistration path was unreachable). These tests pin the per-instance +// semantics and both teardown orders; run them under ASAN/LSAN to verify +// the holder-leak fix and the owner-destroyed-first ordering. + +#include + +#include +#include +#include +#include +#include +#include + +#include "tent/common/concurrent/thread_local_storage.h" + +namespace mooncake { +namespace tent { + +namespace { +struct Cache { + int value = 0; +}; + +size_t countValues(ThreadLocalStorage& storage) { + size_t n = 0; + storage.forEach([&](Cache&) { n++; }); + return n; +} +} // namespace + +// The original aliasing bug: two instances in the same thread must own +// distinct per-thread values. +TEST(ThreadLocalStorageTest, InstancesDoNotAlias) { + ThreadLocalStorage a; + ThreadLocalStorage b; + a.get().value = 42; + EXPECT_EQ(b.get().value, 0); + EXPECT_NE(&a.get(), &b.get()); + b.get().value = 7; + EXPECT_EQ(a.get().value, 42); +} + +// Thread exits while the owner is alive: the value deregisters (forEach no +// longer sees it) and its memory is reclaimed (LSAN would flag the previous +// implementation, which leaked one holder per thread per instantiation). +TEST(ThreadLocalStorageTest, ThreadExitDeregistersAndReclaims) { + ThreadLocalStorage storage; + storage.get().value = 1; // main thread's value + EXPECT_EQ(countValues(storage), 1u); + + std::thread t([&] { + storage.get().value = 2; + EXPECT_EQ(countValues(storage), 2u); + }); + t.join(); + + EXPECT_EQ(countValues(storage), 1u); // worker's value deregistered + EXPECT_EQ(storage.get().value, 1); // main's value untouched +} + +// Owner destroyed while a using thread is still alive: the thread's later +// exit must not touch the dead owner (the jointly-owned control block keeps +// the registry memory valid; ASAN pins this ordering). +TEST(ThreadLocalStorageTest, OwnerDestroyedBeforeThreadExitIsSafe) { + std::atomic used{false}; + std::atomic release{false}; + auto storage = std::make_unique>(); + std::thread t([&] { + storage->get().value = 7; + used.store(true); + while (!release.load()) std::this_thread::yield(); + // Thread exit here runs the node destructor against a dead owner. + }); + while (!used.load()) std::this_thread::yield(); + storage.reset(); // owner gone first + release.store(true); + t.join(); +} + +// Instance ids are never reused: a new storage that may occupy the same +// address as a destroyed one must not see the old orphaned value. +TEST(ThreadLocalStorageTest, DestroyedInstanceStateIsNotResurrected) { + for (int round = 0; round < 8; ++round) { + auto storage = std::make_unique>(); + EXPECT_EQ(storage->get().value, 0) << "round " << round; + storage->get().value = 100 + round; + } +} + +// forEach synchronizes registry membership: it sees exactly the values of +// live threads that used this instance. +TEST(ThreadLocalStorageTest, ForEachVisitsExactlyLiveRegisteredValues) { + ThreadLocalStorage storage; + constexpr int kThreads = 8; + std::atomic ready{0}; + std::atomic release{false}; + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + storage.get().value = i + 1; + ready++; + while (!release.load()) std::this_thread::yield(); + }); + } + while (ready.load() < kThreads) std::this_thread::yield(); + + int sum = 0; + size_t n = 0; + storage.forEach([&](Cache& c) { + sum += c.value; + n++; + }); + EXPECT_EQ(n, (size_t)kThreads); + EXPECT_EQ(sum, kThreads * (kThreads + 1) / 2); + + release.store(true); + for (auto& t : threads) t.join(); + EXPECT_EQ(countValues(storage), 0u); +} + +// Concurrent churn: threads exercising get() across shared storages while +// other storages are created/destroyed, with forEach mixed in. Run under +// TSAN/ASAN for the full effect; asserts basic integrity without them. +TEST(ThreadLocalStorageTest, ConcurrentChurnStress) { + constexpr int kThreads = 8; + constexpr int kIterations = 2000; + ThreadLocalStorage shared_a; + ThreadLocalStorage shared_b; + std::atomic failures{0}; + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + for (int iter = 0; iter < kIterations; ++iter) { + shared_a.get().value = i; + shared_b.get().value = -i; + if (shared_a.get().value != i) failures++; + if (shared_b.get().value != -i) failures++; + // Thread-private storages churn creation/destruction. + ThreadLocalStorage ephemeral; + ephemeral.get().value = iter; + if (ephemeral.get().value != iter) failures++; + if (iter % 64 == 0) { + shared_a.forEach([](Cache&) {}); + } + } + }); + } + for (auto& t : threads) t.join(); + EXPECT_EQ(failures.load(), 0u); + EXPECT_EQ(countValues(shared_a), 0u); +} + +namespace { +struct Counted { + static std::atomic live; + int value = 0; + Counted() { live.fetch_add(1, std::memory_order_relaxed); } + ~Counted() { live.fetch_sub(1, std::memory_order_relaxed); } +}; +std::atomic Counted::live{0}; +} // namespace + +// Storage churn on a long-lived thread: values of destroyed storages must be +// swept by the next first-use get() rather than accumulating until thread +// exit (one orphaned T — pinning its contents — per destroyed storage). +TEST(ThreadLocalStorageTest, OrphanedValuesSweptOnInstanceChurn) { + constexpr int kRounds = 1000; + for (int i = 0; i < kRounds; ++i) { + ThreadLocalStorage storage; + storage.get().value = i; + // Previous round's orphan must have been swept by this round's + // first-use get(): at most this round's value plus one not-yet-swept + // orphan may be alive. + ASSERT_LE(Counted::live.load(), 2) << "round " << i; + } +} + +// Informational: hot-path cost of get(). The remote-desc cache consults this +// on every transfer submit, so the common case must stay a thread_local +// access plus a compare. +TEST(ThreadLocalStorageTest, HotPathMicrobench) { + ThreadLocalStorage storage; + storage.get().value = 1; + constexpr uint64_t kOps = 20'000'000; + volatile int sink = 0; + auto t0 = std::chrono::steady_clock::now(); + for (uint64_t i = 0; i < kOps; ++i) { + sink += storage.get().value; + } + auto t1 = std::chrono::steady_clock::now(); + double ns = + (double)std::chrono::duration_cast(t1 - t0) + .count() / + (double)kOps; + printf("get_hot_path_ns_per_op %.2f\n", ns); + (void)sink; +#if defined(__SANITIZE_THREAD__) || defined(__SANITIZE_ADDRESS__) + constexpr bool kSanitized = true; +#elif defined(__has_feature) +#if __has_feature(thread_sanitizer) || __has_feature(address_sanitizer) + constexpr bool kSanitized = true; +#else + constexpr bool kSanitized = false; +#endif +#else + constexpr bool kSanitized = false; +#endif + // Wall-clock assertions flake under sanitizers (TSAN alone is ~14x); + // elsewhere keep a loose ceiling that still catches syscall- or + // contention-class regressions on the hot path. + if (!kSanitized) { + EXPECT_LT(ns, 100.0); + } +} + +} // namespace tent +} // namespace mooncake From 5507b5cabd74a75b18a8e43fee5a04e3aa5af4d7 Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Mon, 13 Jul 2026 08:21:44 +0800 Subject: [PATCH 53/57] [EP] Cap active RoCE QPs for IBGDA kernels (#2544) Co-authored-by: KMSorSMS Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- mooncake-ep/include/mooncake_ep_api.cuh | 4 +- mooncake-ep/include/mooncake_ep_buffer.h | 5 +++ mooncake-ep/src/mooncake_ep_buffer.cpp | 37 ++++++++++++++++- mooncake-ep/src/mooncake_ep_kernel.cu | 51 ++++++++++++++++++------ 4 files changed, 80 insertions(+), 17 deletions(-) diff --git a/mooncake-ep/include/mooncake_ep_api.cuh b/mooncake-ep/include/mooncake_ep_api.cuh index 2ab1c632a4..1a560d1b90 100644 --- a/mooncake-ep/include/mooncake_ep_api.cuh +++ b/mooncake-ep/include/mooncake_ep_api.cuh @@ -16,7 +16,7 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8, void* workspace, cudaStream_t stream, int64_t timeout_ticks, - int phases); + int phases, int active_qps_per_rank); void mark_phase_ack(void* mxa_buffer, const int32_t* nvlink_available, void* const* ipc_peer_ptrs, int* ack_buffer, int rank, @@ -42,6 +42,6 @@ void combine(void* combined_x, int32_t* active_ranks, void* mxa_buffer, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases, - bool zero_copy); + bool zero_copy, int active_qps_per_rank); } // namespace mooncake diff --git a/mooncake-ep/include/mooncake_ep_buffer.h b/mooncake-ep/include/mooncake_ep_buffer.h index 6088ce20af..03873c5ca2 100644 --- a/mooncake-ep/include/mooncake_ep_buffer.h +++ b/mooncake-ep/include/mooncake_ep_buffer.h @@ -89,6 +89,11 @@ struct MooncakeEpBuffer { bool ibgda_disabled_ = false; int USE_QP_COUNT = MAX_QP_COUNT; + // Cap on active RoCE QPs per peer: spreading small EP messages across too + // many QP/doorbell/progress streams hurts when GPUs share an HCA. Default + // 8; override at runtime with MOONCAKE_EP_ACTIVE_QPS_PER_RANK (>= per-rank + // QP count disables). + int active_qps_cap_ = 8; // Stream for communication at::cuda::CUDAStream comm_stream; diff --git a/mooncake-ep/src/mooncake_ep_buffer.cpp b/mooncake-ep/src/mooncake_ep_buffer.cpp index 79aa4fbe7d..b587e29103 100644 --- a/mooncake-ep/src/mooncake_ep_buffer.cpp +++ b/mooncake-ep/src/mooncake_ep_buffer.cpp @@ -1,11 +1,21 @@ #include #include +#include #include #include #include namespace mooncake { +namespace { + +int active_qps_per_rank_for_ep(int qps_per_rank, bool is_roce, int cap) { + if (!is_roce) return qps_per_rank; + return std::min(qps_per_rank, cap); +} + +} // namespace + // Initialize an RDMA transport: register memory, allocate control buffer, // create QPs. Returns true on success, false if IBGDA is unavailable. static bool initRdmaTransport(device::RdmaTransport* t, void* gdr_buffer, @@ -40,6 +50,23 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks, num_ep_buffer_bytes(num_ep_buffer_bytes), comm_stream(at::cuda::getStreamFromPool(true)) { USE_QP_COUNT = MAX_QP_COUNT / num_ranks * num_ranks; + + // Optional runtime override for the RoCE active-QP cap (default 8). + // Set MOONCAKE_EP_ACTIVE_QPS_PER_RANK to a value >= the per-rank QP count + // (e.g. 256) to effectively disable the cap. + if (const char* env = std::getenv("MOONCAKE_EP_ACTIVE_QPS_PER_RANK")) { + char* end = nullptr; + long v = std::strtol(env, &end, 10); + if (end != env && *end == '\0' && v > 0) { + active_qps_cap_ = static_cast(v); + } else { + LOG(WARNING) << "[EP] ignoring invalid " + "MOONCAKE_EP_ACTIVE_QPS_PER_RANK='" + << env << "'"; + } + } + LOG(INFO) << "[EP] RoCE active QPs/rank cap = " << active_qps_cap_; + // Get ranks CUDA_CHECK(cudaGetDevice(&device_id)); CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate, @@ -209,6 +236,9 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr; int32_t* nvlink_avail = p2p_transport_->availableTablePtr(); void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr(); + int active_qps_per_rank = active_qps_per_rank_for_ep( + USE_QP_COUNT / num_ranks, rdma_transport_ && rdma_transport_->isRoce(), + active_qps_cap_); auto mark_send_done = [=]() { #ifdef MOONCAKE_EP_SPLIT_SEND_RECV @@ -247,7 +277,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, topk_idx.data_ptr(), next_buffer.rdma_recv_signal_buffer, num_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, use_fp8, workspace, launch_stream, - timeout_ticks, phases); + timeout_ticks, phases, active_qps_per_rank); }; if (return_recv_hook) { launcher(LOW_LATENCY_SEND_PHASE); @@ -367,6 +397,9 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr; int32_t* nvlink_avail = p2p_transport_->availableTablePtr(); void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr(); + int active_qps_per_rank = active_qps_per_rank_for_ep( + USE_QP_COUNT / num_ranks, rdma_transport_ && rdma_transport_->isRoce(), + active_qps_cap_); auto mark_send_done = [=]() { #ifdef MOONCAKE_EP_SPLIT_SEND_RECV @@ -405,7 +438,7 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, next_buffer.rdma_recv_signal_buffer, num_combined_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, workspace, launch_stream, timeout_ticks, phases, - zero_copy); + zero_copy, active_qps_per_rank); }; if (return_recv_hook) { launcher(LOW_LATENCY_SEND_PHASE); diff --git a/mooncake-ep/src/mooncake_ep_kernel.cu b/mooncake-ep/src/mooncake_ep_kernel.cu index 0625490b99..bc96441181 100644 --- a/mooncake-ep/src/mooncake_ep_kernel.cu +++ b/mooncake-ep/src/mooncake_ep_kernel.cu @@ -28,6 +28,15 @@ using mooncake::device::mc_atomic_add_release; using mooncake::device::mc_fence; using mooncake::device::mc_fence_barrier_fence; +__device__ __forceinline__ int ep_qp_channel(int expert_local_idx, + int qps_per_rank, + int active_qps_per_rank) { + int active_qps = active_qps_per_rank; + if (active_qps <= 0 || active_qps > qps_per_rank) + active_qps = qps_per_rank; + return expert_local_idx % active_qps; +} + __global__ void mark_phase_ack_kernel(void* mxa_buffer, const int32_t* nvlink_available, void* const* ipc_peer_ptrs, @@ -143,7 +152,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, int num_tokens, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, int64_t timeout_ticks, - int phases) { + int phases, int active_qps_per_rank) { const auto sm_id = static_cast(blockIdx.x); const auto thread_id = static_cast(threadIdx.x); const auto warp_id = thread_id / 32, lane_id = get_lane_id(); @@ -283,8 +292,12 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, mc_fence(); } else { // IBGDA path — send directly from source buffer - mc_rdma_put(comm_ctx, dst_expert_local_idx % num_qp_per_rank, dst_rank, num_qp_per_rank, - src_ptr, dst_ptr, num_bytes_per_msg, lane_id); + mc_rdma_put(comm_ctx, + ep_qp_channel(dst_expert_local_idx, + num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, src_ptr, dst_ptr, + num_bytes_per_msg, lane_id); } // Increase counter after finishing @@ -352,8 +365,11 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, while (mc_ld_acquire(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2); if (dst_rank != rank) { int* signal_ptr = rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank; - mc_red_add(comm_ctx, dst_rank, dst_expert_local_idx % num_qp_per_rank, num_qp_per_rank, - signal_ptr, static_cast(-num_tokens_sent - 1)); + mc_red_add(comm_ctx, dst_rank, + ep_qp_channel(dst_expert_local_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, + static_cast(-num_tokens_sent - 1)); } else { mc_st_release(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank, -num_tokens_sent - 1); } @@ -465,7 +481,8 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, int* next_clean_buffer, int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8, - void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases) { + void* workspace, cudaStream_t stream, int64_t timeout_ticks, + int phases, int active_qps_per_rank) { constexpr int kNumMaxTopK = 11; constexpr int kNumWarpsPerGroup = 4; #ifdef MOONCAKE_EP_USE_MUSA @@ -502,7 +519,8 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \ atomic_counter_per_expert, atomic_finish_counter_per_expert, \ next_clean_buffer, \ num_tokens, num_max_dispatch_tokens_per_rank, \ - num_topk, num_experts, rank, num_ranks, timeout_ticks, phases); } break + num_topk, num_experts, rank, num_ranks, timeout_ticks, phases, \ + active_qps_per_rank); } break SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream); SWITCH_HIDDEN(DISPATCH_LAUNCH_CASE); @@ -526,7 +544,7 @@ combine(void* combined_x, int32_t* active_ranks, int num_max_dispatch_tokens_per_rank, int num_experts, int rank, int num_ranks, int64_t timeout_ticks, - int phases, bool zero_copy) { + int phases, bool zero_copy, int active_qps_per_rank) { const auto sm_id = static_cast(blockIdx.x); const auto num_sms = static_cast(gridDim.x); const auto thread_id = static_cast(threadIdx.x); @@ -610,8 +628,11 @@ combine(void* combined_x, int32_t* active_ranks, if (not zero_copy) UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, buf_int4_ptr, x_int4, mc_ld_nc, mc_st_na); __syncwarp(); - mc_rdma_put(comm_ctx, local_expert_idx % num_qp_per_rank, dst_rank, num_qp_per_rank, - buf_ptr, dst_ptr, num_bytes_per_slot, lane_id); + mc_rdma_put(comm_ctx, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, buf_ptr, dst_ptr, + num_bytes_per_slot, lane_id); } } // Put finishing flag @@ -621,7 +642,10 @@ combine(void* combined_x, int32_t* active_ranks, while (mc_ld_acquire(atomic_clean_flag) == 0); if (dst_rank != rank) { int* signal_ptr = rdma_recv_signal_buffer + global_expert_idx; - mc_signal(comm_ctx, dst_rank, local_expert_idx % num_qp_per_rank, num_qp_per_rank, signal_ptr, 1); + mc_signal(comm_ctx, dst_rank, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, 1); } else { mc_st_release(rdma_recv_signal_buffer + global_expert_idx, 1); } @@ -722,7 +746,8 @@ void combine(void* combined_x, int32_t* active_ranks, int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, void* workspace, cudaStream_t stream, - int64_t timeout_ticks, int phases, bool zero_copy) { + int64_t timeout_ticks, int phases, bool zero_copy, + int active_qps_per_rank) { constexpr int kNumWarpsPerGroup = 4; constexpr int kNumWarpGroups = 8; constexpr int kNumMaxTopk = 11; @@ -751,7 +776,7 @@ LAUNCH_KERNEL(&cfg, combine_func, \ num_combined_tokens, hidden, num_topk, \ num_max_dispatch_tokens_per_rank, \ num_experts, rank, num_ranks, \ - timeout_ticks, phases, zero_copy); } break + timeout_ticks, phases, zero_copy, active_qps_per_rank); } break SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream); SWITCH_HIDDEN(COMBINE_LAUNCH_CASE); From 4293dce8b59132b6bbc77022843a3f85f2888e47 Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Mon, 13 Jul 2026 09:02:38 +0800 Subject: [PATCH 54/57] [EP] add DeepEP V2 elastic buffer (#2503) --- .typos.toml | 6 +- mooncake-ep/benchmarks/elastic_buffer_perf.py | 307 ++++++ .../elastic/mooncake_ep_elastic_api.cuh | 20 + .../elastic/mooncake_ep_elastic_buffer.h | 170 +++ .../mooncake_ep_elastic_combine_official.cuh | 365 +++++++ ...ake_ep_elastic_combine_reduce_epilogue.cuh | 212 ++++ .../mooncake_ep_elastic_combine_utils.cuh | 209 ++++ .../elastic/mooncake_ep_elastic_comm.cuh | 190 ++++ .../elastic/mooncake_ep_elastic_compiled.cuh | 115 +++ ...cake_ep_elastic_dispatch_copy_epilogue.cuh | 277 +++++ ...lastic_dispatch_deterministic_prologue.cuh | 171 +++ .../mooncake_ep_elastic_dispatch_official.cuh | 512 +++++++++ .../elastic/mooncake_ep_elastic_exception.cuh | 81 ++ ...ake_ep_elastic_hybrid_combine_official.cuh | 787 ++++++++++++++ ...ke_ep_elastic_hybrid_dispatch_official.cuh | 888 ++++++++++++++++ .../elastic/mooncake_ep_elastic_launch.cuh | 78 ++ .../elastic/mooncake_ep_elastic_layout.cuh | 374 +++++++ .../elastic/mooncake_ep_elastic_math.cuh | 83 ++ .../elastic/mooncake_ep_elastic_ptx.cuh | 735 +++++++++++++ .../elastic/mooncake_ep_elastic_transport.cuh | 261 +++++ mooncake-ep/include/mooncake_ep_buffer.h | 3 + mooncake-ep/include/mooncake_ep_device.h | 9 +- mooncake-ep/include/mooncake_ep_exception.cuh | 3 + mooncake-ep/setup.py | 7 +- mooncake-ep/src/CMakeLists.txt | 2 +- mooncake-ep/src/ep_py.cpp | 90 ++ .../src/mooncake_ep_elastic_buffer.cpp | 632 ++++++++++++ mooncake-ep/src/mooncake_ep_elastic_kernel.cu | 973 ++++++++++++++++++ mooncake-ep/tests/test_elastic_buffer.py | 412 ++++++++ mooncake-integration/CMakeLists.txt | 1 + .../include/transport/device/p2p_device.cuh | 4 + .../mooncake/mooncake_elastic_buffer.py | 582 +++++++++++ 32 files changed, 8554 insertions(+), 5 deletions(-) create mode 100644 mooncake-ep/benchmarks/elastic_buffer_perf.py create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh create mode 100644 mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh create mode 100644 mooncake-ep/src/mooncake_ep_elastic_buffer.cpp create mode 100644 mooncake-ep/src/mooncake_ep_elastic_kernel.cu create mode 100644 mooncake-ep/tests/test_elastic_buffer.py create mode 100644 mooncake-wheel/mooncake/mooncake_elastic_buffer.py diff --git a/.typos.toml b/.typos.toml index 320972c8fa..4843e506f2 100644 --- a/.typos.toml +++ b/.typos.toml @@ -1,5 +1,5 @@ [default] -extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa"] +extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa", "ue"] [default.extend-words] CANN = "CANN" @@ -9,9 +9,13 @@ wqs = "wqs" # AMD HSA runtime symbol prefix (hsa_*, hsaRes, hsaErr, etc.) — used by the # ROCm dmabuf MR registration path. hsa = "hsa" +Optin = "Optin" HPE = "HPE" [files] extend-exclude = [ "mooncake-transfer-engine/tent/include/tent/thirdparty/nlohmann/json.h", + # DeepEP-derived elastic kernel headers keep upstream identifiers such as + # `ue8m0x4`; exclude the imported header block from spelling checks. + "mooncake-ep/include/elastic/*", ] diff --git a/mooncake-ep/benchmarks/elastic_buffer_perf.py b/mooncake-ep/benchmarks/elastic_buffer_perf.py new file mode 100644 index 0000000000..0bd2d10d13 --- /dev/null +++ b/mooncake-ep/benchmarks/elastic_buffer_perf.py @@ -0,0 +1,307 @@ +#!/usr/bin/env python3 +"""Performance smoke for Mooncake ElasticBuffer dispatch/combine. + +The benchmark intentionally keeps the workload simple and reproducible. It is +not a full system benchmark; it provides a reviewer-friendly way to verify that +the new elastic path runs repeatedly, supports cached handles, and reports +per-rank effective payload bandwidth. + +Typical single-node usage: + + MOONCAKE_EP_NUM_LOCAL_RANKS=8 \ + torchrun --standalone --nproc_per_node=8 \ + mooncake-ep/benchmarks/elastic_buffer_perf.py --route alltoall +""" + +from __future__ import annotations + +import argparse +import os +import time +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.testing as testing + +from mooncake.mooncake_elastic_buffer import ElasticBuffer + + +@dataclass(frozen=True) +class RoutePlan: + topk_idx: torch.Tensor + expected_recv_tokens: int + expected_combine_factor: int + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark Mooncake ElasticBuffer") + parser.add_argument("--num-tokens", type=int, default=128) + parser.add_argument("--max-tokens", type=int, default=0) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--num-topk", type=int, default=8) + parser.add_argument("--num-sms", type=int, default=24) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument( + "--route", + choices=("alltoall", "local", "cross"), + default="alltoall", + help="Expert routing pattern to generate.", + ) + parser.add_argument( + "--reuse-handle", + action=argparse.BooleanOptionalAction, + default=True, + help="Reuse the first dispatch handle for later iterations.", + ) + parser.add_argument( + "--check-correctness", + action=argparse.BooleanOptionalAction, + default=True, + help="Check combine output on each iteration.", + ) + parser.add_argument( + "--sync-actual-count", + action="store_true", + help="Synchronize and verify GPU-side received-token count each iteration.", + ) + parser.add_argument("--seed", type=int, default=2026) + return parser.parse_args() + + +def init_distributed(seed: int) -> tuple[int, int]: + if not dist.is_initialized(): + dist.init_process_group("nccl") + rank = dist.get_rank() + local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count())) + torch.cuda.set_device(local_rank) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(seed + rank) + return rank, dist.get_world_size() + + +def make_route_plan( + *, + rank: int, + world_size: int, + buffer: ElasticBuffer, + num_tokens: int, + num_topk: int, + num_experts: int, + route: str, +) -> RoutePlan: + local_experts = num_experts // world_size + if local_experts <= 0: + raise ValueError("num_experts must be at least world_size") + expert_offsets = torch.arange(num_topk, device="cuda", dtype=torch.long) % local_experts + + if route == "cross" and buffer.num_scaleout_ranks > 1: + dst_scaleout = (buffer.scaleout_rank_idx + 1) % buffer.num_scaleout_ranks + dst_rank = dst_scaleout * buffer.num_scaleup_ranks + buffer.scaleup_rank_idx + choices = dst_rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + if route == "local" or (route == "cross" and buffer.num_scaleout_ranks == 1): + choices = rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + dst_ranks = (rank + torch.arange(num_topk, device="cuda", dtype=torch.long)) % world_size + choices = dst_ranks * local_experts + expert_offsets + unique_dst_ranks = int(torch.unique(dst_ranks).numel()) + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens * unique_dst_ranks, + unique_dst_ranks, + ) + + +def make_input(rank: int, iteration: int, num_tokens: int, hidden: int) -> torch.Tensor: + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + return (base + rank * 1_000_000 + iteration * 17).to(torch.bfloat16).contiguous() + + +def check_output( + *, + rank: int, + route: str, + combined: torch.Tensor, + expected: torch.Tensor, +) -> None: + if route == "local": + if not torch.equal(combined, expected): + diff = (combined.float() - expected.float()).abs().max().item() + raise AssertionError(f"rank={rank}: local-route mismatch, max_diff={diff}") + return + + testing.assert_close( + combined, + expected, + rtol=1e-2, + atol=1e-3, + msg=lambda msg: f"rank={rank}: {route} combine mismatch: {msg}", + ) + + +def main() -> None: + args = parse_args() + rank, world_size = init_distributed(args.seed) + max_tokens = args.max_tokens or max(128, args.num_tokens) + num_experts = args.num_experts + if num_experts % world_size != 0: + raise ValueError("num_experts must be divisible by world_size") + + buffer = ElasticBuffer( + dist.group.WORLD, + num_max_tokens_per_rank=max_tokens, + hidden=args.hidden, + num_topk=args.num_topk, + use_fp8_dispatch=False, + deterministic=False, + allow_hybrid_mode=True, + allow_multiple_reduction=True, + num_gpu_timeout_secs=10, + ) + route_plan = make_route_plan( + rank=rank, + world_size=world_size, + buffer=buffer, + num_tokens=args.num_tokens, + num_topk=args.num_topk, + num_experts=num_experts, + route=args.route, + ) + weights = torch.ones((args.num_tokens, args.num_topk), device="cuda", dtype=torch.float32) + + def run_one(iteration: int, cached_handle): + x = make_input(rank, iteration, args.num_tokens, args.hidden) + dispatch_start = torch.cuda.Event(enable_timing=True) + dispatch_end = torch.cuda.Event(enable_timing=True) + combine_end = torch.cuda.Event(enable_timing=True) + + use_cached = args.reuse_handle and cached_handle is not None + dispatch_start.record() + recv_x, _recv_idx, recv_weights, handle, _ = buffer.dispatch( + x, + topk_idx=None if use_cached else route_plan.topk_idx, + topk_weights=None if use_cached else weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + handle=cached_handle if use_cached else None, + do_cpu_sync=True if not use_cached else None, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + dispatch_end.record() + + actual_recv_tokens = route_plan.expected_recv_tokens + if args.sync_actual_count: + actual_recv_tokens = int(handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if actual_recv_tokens != route_plan.expected_recv_tokens: + raise AssertionError( + f"rank={rank}: got {actual_recv_tokens} received tokens, " + f"expected {route_plan.expected_recv_tokens}" + ) + + combined, _combined_weights, _ = buffer.combine( + recv_x[:actual_recv_tokens].contiguous(), + handle, + topk_weights=( + recv_weights[:actual_recv_tokens].contiguous() + if recv_weights is not None + else None + ), + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + combine_end.record() + torch.cuda.synchronize() + + if args.check_correctness: + expected = (x.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_output(rank=rank, route=args.route, combined=combined, expected=expected) + + return ( + handle, + dispatch_start.elapsed_time(dispatch_end), + dispatch_end.elapsed_time(combine_end), + actual_recv_tokens, + ) + + cached_handle = None + for i in range(args.warmup): + cached_handle, _dispatch_ms, _combine_ms, _actual = run_one(i, cached_handle) + + dist.barrier() + torch.cuda.synchronize() + dispatch_ms = [] + combine_ms = [] + recv_tokens = [] + wall_start = time.time() + for i in range(args.iters): + cached_handle, d_ms, c_ms, actual = run_one(args.warmup + i, cached_handle) + dispatch_ms.append(d_ms) + combine_ms.append(c_ms) + recv_tokens.append(actual) + torch.cuda.synchronize() + dist.barrier() + wall_seconds = time.time() - wall_start + + stats = torch.tensor( + [ + sum(dispatch_ms) / len(dispatch_ms), + sum(combine_ms) / len(combine_ms), + min(dispatch_ms), + max(dispatch_ms), + min(combine_ms), + max(combine_ms), + sum(recv_tokens) / len(recv_tokens), + wall_seconds, + ], + device="cuda", + dtype=torch.float64, + ) + gathered = [torch.empty_like(stats) for _ in range(world_size)] + dist.all_gather(gathered, stats) + + if rank == 0: + table = torch.stack(gathered).cpu() + payload_bytes = table[:, 6].mean().item() * args.hidden * 2 + dispatch_avg_ms = table[:, 0].mean().item() + combine_avg_ms = table[:, 1].mean().item() + print( + "MOONCAKE_ELASTIC_PERF_OK", + f"world={world_size}", + f"route={args.route}", + f"reuse_handle={int(args.reuse_handle)}", + f"tokens={args.num_tokens}", + f"hidden={args.hidden}", + f"topk={args.num_topk}", + f"scaleout={buffer.num_scaleout_ranks}", + f"scaleup={buffer.num_scaleup_ranks}", + f"dispatch_avg_ms={dispatch_avg_ms:.3f}", + f"combine_avg_ms={combine_avg_ms:.3f}", + f"recv_tokens_avg={table[:, 6].mean().item():.1f}", + f"effective_payload_MB_per_rank={payload_bytes / 1e6:.1f}", + f"dispatch_effective_GBps={payload_bytes / dispatch_avg_ms / 1e6:.2f}", + f"combine_effective_GBps={payload_bytes / combine_avg_ms / 1e6:.2f}", + flush=True, + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh new file mode 100644 index 0000000000..aa99b8e0c9 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh @@ -0,0 +1,20 @@ +#pragma once + +// Official DeepEP elastic source import surface for Mooncake. +// +// This umbrella intentionally lives under include/elastic, not in the legacy EP +// include root. It keeps the imported elastic implementation discoverable +// while allowing the host launch/runtime glue to opt in file-by-file without +// perturbing legacy Buffer dispatch/combine symbols. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h b/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h new file mode 100644 index 0000000000..6b5dacdd07 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h @@ -0,0 +1,170 @@ +#ifndef MOONCAKE_EP_ELASTIC_BUFFER_H +#define MOONCAKE_EP_ELASTIC_BUFFER_H + +#include +#include +#include +#include +#include +#include + +#include + +namespace mooncake { + +struct ElasticLaunchContext; + +struct ElasticTopology { + int rank_idx = 0; + int num_ranks = 1; + int num_rdma_ranks = 1; + int num_nvlink_ranks = 1; + int num_scaleout_ranks = 1; + int num_scaleup_ranks = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + bool hybrid_enabled = false; +}; + +struct ElasticConfig { + int64_t num_max_tokens_per_rank = 0; + int64_t hidden = 0; + int64_t num_topk = 0; + bool use_fp8_dispatch = false; + bool deterministic = false; + bool allow_hybrid_mode = true; + bool allow_multiple_reduction = true; + bool prefer_overlap_with_compute = true; + int sl_idx = 3; + int num_allocated_qps = 0; + int num_cpu_timeout_secs = 300; + int num_gpu_timeout_secs = 100; +}; + +struct ElasticNativeHandle { + bool do_expand = false; + int num_experts = 0; + int expert_alignment = 1; + int num_max_tokens_per_rank = 0; + int num_sms = 0; + torch::Tensor topk_idx; + torch::Tensor psum_num_recv_tokens_per_scaleup_rank; + torch::Tensor psum_num_recv_tokens_per_expert; + torch::Tensor recv_src_metadata; + torch::Tensor recv_layout_range; + torch::Tensor dst_buffer_slot_idx; + std::optional token_metadata_at_forward; + std::optional channel_linked_list; + std::vector num_recv_tokens_per_expert_list; +}; + +struct ElasticDispatchOutput { + torch::Tensor recv_x; + std::optional recv_x_scales; + std::optional recv_topk_idx; + std::optional recv_topk_weights; + ElasticNativeHandle handle; + std::optional event; +}; + +struct ElasticCombineOutput { + torch::Tensor combined_x; + std::optional combined_topk_weights; + std::optional event; +}; + +class MooncakeElasticBuffer { + public: + MooncakeElasticBuffer(int rank, int num_ranks, int64_t num_buffer_bytes, + int64_t num_max_tokens_per_rank, int64_t hidden, + int64_t num_topk, bool use_fp8_dispatch, + bool deterministic, bool allow_hybrid_mode, + bool allow_multiple_reduction, + bool prefer_overlap_with_compute, int sl_idx, + int num_allocated_qps, int num_cpu_timeout_secs, + int num_gpu_timeout_secs); + + ~MooncakeElasticBuffer(); + + static int64_t calculate_buffer_size(int num_ranks, + int64_t num_max_tokens_per_rank, + int64_t hidden, int64_t num_topk, + bool use_fp8_dispatch, + bool allow_hybrid_mode, + bool allow_multiple_reduction); + + std::tuple get_physical_domain_size() const; + std::tuple get_logical_domain_size() const; + int get_theoretical_num_sms(int num_experts, int num_topk) const; + + ElasticDispatchOutput dispatch( + const torch::Tensor& x, const std::optional& sf, + const torch::Tensor& topk_idx, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_experts, + int num_max_tokens_per_rank, int expert_alignment, int num_sms, + bool do_expand, bool do_cpu_sync, bool async_with_compute_stream, + const std::optional& cached_handle = std::nullopt); + + ElasticCombineOutput combine( + const torch::Tensor& x, const ElasticNativeHandle& handle, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_sms, + bool async_with_compute_stream, + const std::optional& out); + + MooncakeEpBuffer& native_buffer() { return *native_buffer_; } + + bool ibgda_disabled() const { return native_buffer_->ibgda_disabled(); } + bool use_fast_path() { return native_buffer_->use_fast_path(); } + void update_local_qpns() { native_buffer_->update_local_qpns(); } + bool is_roce() const { return native_buffer_->is_roce(); } + void sync_ibgda_peers(const std::vector& remote_addrs, + const std::vector& remote_keys, + const std::vector>& peer_qpns, + const std::vector>& peer_lids, + const std::vector& subnet_prefixes, + const std::vector& interface_ids, + const std::vector& active_ranks_mask) { + native_buffer_->sync_ibgda_peers(remote_addrs, remote_keys, peer_qpns, + peer_lids, subnet_prefixes, + interface_ids, active_ranks_mask); + } + std::tuple get_mr_info() { + return native_buffer_->get_mr_info(); + } + std::tuple get_gid() { return native_buffer_->get_gid(); } + std::vector get_local_qpns() { + return native_buffer_->get_local_qpns(); + } + std::vector get_local_lids() { + return native_buffer_->get_local_lids(); + } + std::vector get_ipc_handle() { + return native_buffer_->get_ipc_handle(); + } + void sync_nvlink_ipc_handles( + const std::vector>& remote_handles, + const std::vector& active_ranks_mask) { + native_buffer_->sync_nvlink_ipc_handles(remote_handles, + active_ranks_mask); + } + + private: + ElasticConfig config_; + ElasticTopology topology_; + std::unique_ptr native_buffer_; + int64_t host_workspace_bytes_ = 0; + void* host_workspace_ = nullptr; + void* mapped_host_workspace_ = nullptr; + + static ElasticLaunchContext make_launch_context( + MooncakeEpBuffer& buffer, const ElasticTopology& topology, + void* mapped_host_workspace, int64_t timeout_cycles); + static ElasticTopology discover_topology(int rank, int num_ranks, + bool allow_hybrid_mode); +}; + +} // namespace mooncake + +#endif // MOONCAKE_EP_ELASTIC_BUFFER_H diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh new file mode 100644 index 0000000000..9e36b344f7 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh @@ -0,0 +1,365 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include + +#include + +namespace mooncake::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumRanks, kNumTopk>(), + typename team_t = std::conditional_t< + kIsScaleupNVLink, transport::ScaleupTeam, transport::WorldTeam>> +__global__ void __launch_bounds__(kNumThreads, 1) + combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + const device::CommCtx comm_ctx, void* buffer, void* workspace, + const int rank_idx, int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = (ptx::get_warp_idx() + rank_idx) % kNumWarps; + const auto lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + constexpr bool kDoExpandedSend = + not kAllowMultipleReduction and kUseExpandedLayout; + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = + __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumRanks - 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + const auto recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, buffer); + const auto send_buffer = layout::BufferLayout( + token_layout, kNumRanks, + kNumMaxTokensPerRank * (kDoExpandedSend ? kNumTopk : 1), + recv_buffer.get_buffer_end_ptr()); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode(sm_idx, warp_idx); + const auto gin = transport::MooncakeGin(comm_ctx, qp_idx, sharing_mode, + kNumQPs, 0, 0, 0, kNumRanks); + + // Full barrier to ensure the remote buffer is available + const auto workspace_layout = + layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Do TMA writes into the remote buffers + int num_tokens_per_warp = + math::ceil_div(num_reduced_tokens, kNumSMs * kNumWarps); + const int token_start_idx = num_tokens_per_warp * global_warp_idx; + const int token_end_idx = + min(token_start_idx + num_tokens_per_warp, num_reduced_tokens); + for (int i = token_start_idx; i < token_end_idx; ++i) { + // The master slot index during dispatch + constexpr int kMetadataStride = 2 + kNumTopk; + const int src_token_idx = + __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank; + const int src_rank_topk_idx = + __ldg(src_metadata + i * kMetadataStride + 1); + const int src_rank_idx = src_rank_topk_idx / kNumTopk; + const int src_topk_idx = src_rank_topk_idx % kNumTopk; + + // Directly to the remote or via RDMA + const bool nvlink_bypass = + gin.is_nvlink_accessible(src_rank_idx); + layout::TokenLayout master_token_buffer = [=]() { + // NVLink bypass + if (nvlink_bypass) { + auto token_buffer = + recv_buffer + .get_rank_buffer(kUseRankLayout ? rank_idx + : src_topk_idx) + .get_token_buffer(src_token_idx); + token_buffer.set_base_ptr(gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx)); + return token_buffer; + } + + // Use RDMA + return send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx); + }(); + + // Hidden requirements + EP_STATIC_ASSERT( + kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, + "Invalid hidden"); + using combine_vec_t = + typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = + kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = + __ldg(src_metadata + i * kMetadataStride + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no expand + no reduce, or expand + no reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = + not kUseExpandedLayout or + (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = i; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = + ptx::exchange(stored_topk_slot_idx, + ptx::get_master_lane_idx(reduce_valid_mask)); + + // No reduce +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto src_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + auto* dst_ptr = static_cast( + master_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, src_ptr[vec_idx]); + } + __syncwarp(); + __threadfence_system(); + } +#else + if (ptx::elect_one_sync()) { + const auto load_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, + mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + }); + + // Reduce into shared memory + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA stores +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto* src_ptr = static_cast( + tma_buffer.get_base_ptr()); + auto* dst_ptr = static_cast( + master_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, src_ptr[vec_idx]); + } + __syncwarp(); + __threadfence_system(); + } +#else + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + } else { +// No local reduction, send all data (expanded send) +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + const auto slot_idx = ptx::exchange(stored_topk_slot_idx, k); + if (slot_idx >= 0) { + const auto src_token_ptr = math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + const auto token_buffer = + recv_buffer.get_rank_buffer(k).get_token_buffer( + src_token_idx); +#ifdef MOONCAKE_EP_USE_MUSA + if (nvlink_bypass) { + auto* dst_ptr = + static_cast(gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx)); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, + src_token_ptr[vec_idx]); + } + } else { + const auto send_token_buffer = + send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx * kNumTopk + k); + auto* dst_ptr = static_cast( + send_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, + src_token_ptr[vec_idx]); + } + __syncwarp(); + if (ptx::elect_one_sync()) { + gin.put(token_buffer.get_base_ptr(), + send_token_buffer.get_base_ptr(), + kNumHiddenBytes, src_rank_idx); + } + } + __syncwarp(); + __threadfence_system(); +#else + if (ptx::elect_one_sync()) { + // Load + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + src_token_ptr, mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + if (nvlink_bypass) { + // Write into the same position + ptx::tma_store_1d( + gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } else { + // Write to the RDMA send buffer + const auto send_token_buffer = + send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx * kNumTopk + + k); + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + kNumHiddenBytes); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue RDMA + gin.put(token_buffer.get_base_ptr(), + send_token_buffer.get_base_ptr(), + kNumHiddenBytes, src_rank_idx); + } + } + __syncwarp(); +#endif + } + } + } + + // Write topk weights + if (not kUseExpandedLayout and topk_weights != nullptr and + lane_idx < kNumTopk) { + const float value = __ldg(topk_weights + (i * kNumTopk + lane_idx)); +#ifdef MOONCAKE_EP_USE_MUSA + ptx::st_relaxed_sys( + master_token_buffer.get_topk_weights_ptr() + lane_idx, value); +#else + master_token_buffer.get_topk_weights_ptr()[lane_idx] = value; +#endif + } + __syncwarp(); +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); +#endif + + // Wait send buffer's TMA store and issue RDMA send + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and not nvlink_bypass and + ptx::elect_one_sync()) { + ptx::tma_store_wait(); + const auto dst_ptr = + recv_buffer + .get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx) + .get_token_buffer(src_token_idx) + .get_base_ptr(); + gin.put(dst_ptr, master_token_buffer.get_base_ptr(), + master_token_buffer.get_num_bytes(), + src_rank_idx); + } + } + + // Final barrier to ensure data arrival + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh new file mode 100644 index 0000000000..6e41aebab8 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh @@ -0,0 +1,212 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include + +#include + +namespace mooncake::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumRanks, kNumTopk>()> +__global__ void __launch_bounds__(kNumThreads, 1) + combine_reduce_epilogue_impl(nv_bfloat16* combined_x, + float* combined_topk_weights, + topk_idx_t* combined_topk_idx, + void* recv_buffer, void* bias_0, void* bias_1, + const int num_combined_tokens, + const int scaleout_rank_idx, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + constexpr int kNumExpertsPerRank = + kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks); + EP_STATIC_ASSERT(kNumExperts % (kNumScaleupRanks * kNumScaleoutRanks) == 0, + "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = + warp_idx * kNumSMs + + sm_idx; // NOTES: Here we prioritize distributing tasks to different + // SMs to ensure that the last wave is evenly concentrated on + // each SM. + + // Load buffers from scale-out or scale-up ranks + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto comm_token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto comm_buffer = + layout::BufferLayout(comm_token_layout, kNumTokensInLayout, + kNumMaxTokensPerRank, recv_buffer); + + // Store buffers + const auto output_token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, 0, false); + const auto output_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, combined_x); + const auto tma_buffer = + layout::BufferLayout(output_token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + + // Bias layout + const auto bias_0_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, bias_0); + const auto bias_1_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, bias_1); + + // Will block until the main combine kernel has finished and all data are + // visible NOTES: PDL is used, please do not use `__ldg` +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // Read from buffers and do reduction + for (int token_idx = global_warp_idx; token_idx < num_combined_tokens; + token_idx += kNumWarps * kNumSMs) { + // Preprocess all indices + int stored_dst_rank_idx = -1, stored_dst_expert_idx = -1; + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + if (lane_idx < kNumTopk) { + stored_dst_expert_idx = static_cast( + combined_topk_idx[token_idx * kNumTopk + lane_idx]); + stored_dst_rank_idx = + stored_dst_expert_idx >= 0 + ? stored_dst_expert_idx / (kNumScaleoutRanks == 1 + ? kNumExpertsPerRank + : kNumExpertsPerScaleout) + : -1; + } + __syncwarp(); + + // Sort valid top-k indices to front + const auto [should_deduplicate, + deduplicate_key] = [&]() -> std::pair { + if constexpr (kUseExpandedLayout and not kAllowMultipleReduction) { + // Activations are never reduced before + return {false, 0}; + } else if constexpr (kNumScaleoutRanks != 1 and + not kUseExpandedLayout and + not kAllowMultipleReduction) { + // Hybrid mode without expanded layout and multiple reduction. + // Should deduplicate on a per-rank basis + return {true, stored_dst_expert_idx >= 0 + ? stored_dst_expert_idx / kNumExpertsPerRank + : -1}; + } else { + // Should deduplicate on a per-rank (for non-hybrid mode) or a + // per-scale-rank (for hybrid mode) basis + return {true, stored_dst_rank_idx}; + } + }(); + auto reduce_valid_mask = + should_deduplicate + ? ptx::gather(ptx::deduplicate(deduplicate_key, lane_idx) and + stored_dst_rank_idx >= 0) + : ptx::gather(stored_dst_rank_idx >= 0); + int topk_slot_idx[kNumTokensInLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return kUseRankLayout ? ptx::exchange(stored_dst_rank_idx, idx) + : idx; + }); + + // Iterate over per-hidden-chunk stage + using combine_vec_t = + typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = + kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return static_cast( + comm_buffer.get_rank_buffer(slot_idx) + .get_token_buffer(token_idx) + .get_base_ptr()); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }, + /* Bias 0 */ bias_0 == nullptr + ? nullptr + : static_cast( + bias_0_buffer.get_token_buffer(token_idx).get_base_ptr()), + /* Bias 1 */ bias_1 == nullptr + ? nullptr + : static_cast( + bias_1_buffer.get_token_buffer(token_idx) + .get_base_ptr())); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA copy +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto* src_ptr = + static_cast(tma_buffer.get_base_ptr()); + auto* dst_ptr = static_cast( + output_buffer.get_token_buffer(token_idx).get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; vec_idx += 32) { + dst_ptr[vec_idx] = src_ptr[vec_idx]; + } + __syncwarp(); + } +#else + if (ptx::elect_one_sync()) { + ptx::tma_store_1d( + output_buffer.get_token_buffer(token_idx).get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + + // Write top-k weights + if (combined_topk_weights != nullptr) { + const auto master_lane_idx = + ptx::get_master_lane_idx(ptx::match(stored_dst_rank_idx)); + if (lane_idx < kNumTopk) { + float value = 0; + if (stored_dst_rank_idx >= 0) { + const auto dst_ptr = + comm_buffer + .get_rank_buffer(kUseRankLayout + ? stored_dst_rank_idx + : master_lane_idx) + .get_token_buffer(token_idx) + .get_topk_weights_ptr() + + lane_idx; + value = *dst_ptr; + } + combined_topk_weights[token_idx * kNumTopk + lane_idx] = value; + } + __syncwarp(); + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh new file mode 100644 index 0000000000..b2e3623f2a --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh @@ -0,0 +1,209 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +namespace mooncake::elastic { + +template +constexpr bool use_rank_layout() { + if constexpr (not kAllowMultipleReduction) return false; + return kNumRanks <= kNumTopk; +} + +template +constexpr int get_num_tokens_in_layout() { + return use_rank_layout() + ? kNumRanks + : kNumTopk; +} + +template +constexpr int get_max_unroll_factor() { + for (int i = kMaxUnrollFactor; i >= 1; --i) + if (kLength % (kWarpSize * i) == 0) return i; +#ifdef MOONCAKE_EP_USE_MUSA + return 1; +#else + throw std::logic_error("Invalid length, cannot find unrolling factor"); +#endif +} + +// Determine the vector type for combine loads/stores based on arch and hidden +// size alignment +template +struct CombineVecTraits { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) + // On SM100+, use longlong4_t (32 bytes) if hidden is aligned, otherwise + // fall back to int4 (16 bytes) + static constexpr bool kUseLonglong4 = + (kHiddenBytes % sizeof(longlong4_t) == 0) and + ((kHiddenBytes / sizeof(longlong4_t)) % 32 == 0); + using vec_t = std::conditional_t; +#else + using vec_t = int4; +#endif +}; + +template +__device__ __forceinline__ void compute_topk_slots( + int (&topk_slot_idx)[kNumValidTopk], uint32_t mask, + const fetch_func_t& fetch_func) { +#pragma unroll + for (int k = 0; k < kNumValidTopk; ++k) { + const int lowest_idx = __ffs(mask) - 1; + // Here we perform the exchange unconditionally to avoid `BRA.DIV` + const auto fetched = fetch_func(lowest_idx); + mask &= mask - 1; + topk_slot_idx[k] = lowest_idx >= 0 ? fetched : -1; + } +} + +template +__device__ __forceinline__ void combine_reduce( + const int& lane_idx, int (&topk_slot_idx)[kNumValidTopk], + vec_t* dst_buffer_ptr, + const get_src_buffer_ptr_func_t& get_src_buffer_ptr_func, + const wait_buffer_func_t& wait_buffer_func, vec_t* bias_0 = nullptr, + vec_t* bias_1 = nullptr) { + constexpr int kNumElemsPerVec = sizeof(vec_t) / sizeof(nv_bfloat16); + EP_STATIC_ASSERT(kNumElemsPerVec % 2 == 0, "Invalid number of elements"); + EP_STATIC_ASSERT(kHiddenVec % (kUnrollFactor * 32) == 0, + "Invalid unrolling"); + + // We use BF16 add as much as possible, as casting is slow + const bool enable_hadd_bypass = + (bias_0 == nullptr and bias_1 == nullptr) and + (kNumValidTopk <= 2 or topk_slot_idx[2] < 0); + EP_STATIC_ASSERT(kNumValidTopk > 0, "Invalid top-k"); + + if (enable_hadd_bypass) { +#pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++i) { + // Read values 0 + const auto slot_0 = topk_slot_idx[0]; + const auto src_base_ptr_0 = get_src_buffer_ptr_func(slot_0); + vec_t values_0[kUnrollFactor] = {}; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values_0[j] = ptx::ldg_with_gez_pred( + src_base_ptr_0 + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + slot_0); + } + + // Read values 1 + vec_t values_1[kUnrollFactor] = {}; + const auto slot_1 = kNumValidTopk == 1 ? -1 : topk_slot_idx[1]; + const auto src_base_ptr_1 = get_src_buffer_ptr_func(slot_1); +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values_1[j] = ptx::ldg_with_gez_pred( + src_base_ptr_1 + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + slot_1); + } + + // Wait buffer releases for the first write + if (i == 0) wait_buffer_func(); + + // Reduce into shared memory + const auto bf162_view_0 = reinterpret_cast(values_0); + const auto bf162_view_1 = reinterpret_cast(values_1); +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { +#pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++l) { + const int idx = j * (kNumElemsPerVec / 2) + l; +#ifdef MOONCAKE_EP_USE_MUSA + bf162_view_0[idx] = __floats2bfloat162_rn( + __low2float(bf162_view_0[idx]) + + __low2float(bf162_view_1[idx]), + __high2float(bf162_view_0[idx]) + + __high2float(bf162_view_1[idx])); +#else + bf162_view_0[idx] += bf162_view_1[idx]; +#endif + } + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = + values_0[j]; + } + } + } else { +#pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++i) { + // Add bias + float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {}; + const auto add_bias = [&](const vec_t* base_ptr) { + // Read + vec_t values[kUnrollFactor]; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) + values[j] = ptx::ldg(base_ptr + i * (kUnrollFactor * 32) + + j * 32 + lane_idx); + + // Reduce + const auto bf162_view = reinterpret_cast(values); +#pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++j) + ptx::accumulate(reduced[j], bf162_view[j]); + }; + bias_0 != nullptr ? add_bias(bias_0) : void(); + bias_1 != nullptr ? add_bias(bias_1) : void(); + +#pragma unroll + for (int k = 0; k < kNumValidTopk; ++k) { + // We have a limitation on `k` to reduce the branch instruction + // count + if (k >= kNumExpectedTopk and topk_slot_idx[k] < 0) break; + + // Read values + const auto src_base_ptr = + get_src_buffer_ptr_func(topk_slot_idx[k]); + vec_t values[kUnrollFactor] = {}; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values[j] = ptx::ldg_with_gez_pred( + src_base_ptr + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + topk_slot_idx[k]); + } + + // Reduce + const auto bf162_view = reinterpret_cast(values); +#pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++j) + ptx::accumulate(reduced[j], bf162_view[j]); + } + + // Wait buffer releases for the first write + if (i == 0) wait_buffer_func(); + +// Cast into shared memory +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + vec_t casted_value; + auto bf162_view = + reinterpret_cast(&casted_value); +#pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++l) { + const auto value = reduced[j * (kNumElemsPerVec / 2) + l]; +#ifdef MOONCAKE_EP_USE_MUSA + bf162_view[l] = __floats2bfloat162_rn(value.x, value.y); +#else + bf162_view[l] = __float22bfloat162_rn(value); +#endif + } + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = + casted_value; + } + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh new file mode 100644 index 0000000000..886c2f12d8 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh @@ -0,0 +1,190 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace mooncake::elastic::comm { + +static constexpr int64_t kNumOneSecCycles = 2000000000; + +static constexpr int kDeviceBarrierTag = 0; +static constexpr int kKernelBarrierTag = 1; +static constexpr int kDispatchTag0 = 2; +static constexpr int kDispatchTag1 = 3; +static constexpr int kCombineTag0 = 4; +static constexpr int kCombineTag1 = 5; +static constexpr int kHybridDispatchTag0 = 6; +static constexpr int kHybridDispatchTag1 = 7; +static constexpr int kHybridCombineTag0 = 8; +static constexpr int kHybridCombineTag1 = 9; + +static constexpr int kFlushAllAllocatedQPs = -1; + +template +__device__ __forceinline__ void timeout_while(const bool& condition, + const func_t& func, + int64_t start_clock = 0) { + if (start_clock == 0) start_clock = clock64(); + while (condition) { + const bool timeout = kNumTimeoutCycles >= 0 && + (clock64() - start_clock >= kNumTimeoutCycles); + if (func(timeout)) break; + if (timeout) { + const auto timeout_start = clock64(); + while (clock64() - timeout_start < kNumOneSecCycles) { + } + ptx::trap(); + } + } +} + +template +__device__ __forceinline__ void timeout_while(const func_t& func, + const int64_t& start_clock = 0) { + timeout_while(true, func, start_clock); +} + +template +__forceinline__ __device__ void local_grid_sync( + const layout::WorkspaceLayout& workspace, const int& thread_idx) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)kNumThreads; + __shared__ unsigned long long ticket; + __syncthreads(); + if (thread_idx == 0) { + ticket = atomicAdd( + workspace.get_nvl_barrier_counter_ptr(kKernelBarrierTag), 1ULL); + } + __syncthreads(); + const auto target = ((ticket / kNumSMs) + 1ULL) * kNumSMs; + timeout_while(thread_idx == 0, [=](const bool&) { + return ptx::ld_volatile( + workspace.get_nvl_barrier_counter_ptr(kKernelBarrierTag)) >= + target; + }); + __syncthreads(); +#else + (void)workspace; + (void)thread_idx; + (gridDim.x > 1) ? cooperative_groups::this_grid().sync() : __syncthreads(); +#endif +} + +template +__device__ __forceinline__ std::pair get_qp_mode( + const int& sm_idx, const int& channel_in_sm_idx, + const bool& is_notify_warp = false) { + if constexpr (kNumQPs == 1) return {0, 1}; + if (is_notify_warp) return {0, 0}; + + constexpr int kQPStartIdx = static_cast(kWithNotifyWarps); + constexpr int kNumAvailableQPs = kNumQPs - kQPStartIdx; + if constexpr (kNumSMs <= kNumAvailableQPs) { + const int num_qps_in_sm = (kNumAvailableQPs / kNumSMs) + + (sm_idx < (kNumAvailableQPs % kNumSMs)); + return {kQPStartIdx + sm_idx + + (channel_in_sm_idx % max(1, num_qps_in_sm)) * kNumSMs, + 0}; + } else { + const auto global_channel_idx = + sm_idx * kNumChannelsPerSM + channel_in_sm_idx; + return {kQPStartIdx + (global_channel_idx % max(1, kNumAvailableQPs)), + 1}; + } +} + +template +__forceinline__ __device__ void mooncake_barrier_wo_local_sync( + const transport::MooncakeGin& gin, const layout::WorkspaceLayout& workspace, + const int& rank_idx, const int& sm_idx, const int& thread_idx) { + if (kNumSMs > 1 && sm_idx > 0) return; + + const int status = + static_cast((*workspace.get_nvl_barrier_counter_ptr(kTag)) & 3); + const int phase = status & 1; + const int sign = status >> 1; + const int* base_signal = workspace.get_nvl_barrier_signal_ptr(kTag, phase); + + if (thread_idx < kNumRanks) { + auto* dst_ptr = const_cast(base_signal) + rank_idx; + gin.red_add_rel(dst_ptr, sign ? -1 : 1, thread_idx); + } + __syncthreads(); + + if (thread_idx == 0) + atomicAdd(workspace.get_nvl_barrier_counter_ptr(kTag), 1ULL); + + timeout_while( + thread_idx == 0, [=](const bool& is_last_check) { + int sum = 0; +#pragma unroll + for (int i = 0; i < kNumRanks; ++i) { + sum += + ptx::ld_acquire_sys(const_cast(base_signal) + i); + } + // Mooncake's portable barrier uses one additive slot per source + // rank. Each positive phase adds +1 into a zeroed phase slot; the + // matching negative phase later adds -1 into the same phase slot. + // This matches RDMA atomic-add semantics and avoids relying on a + // remote store primitive for non-P2P peers. + const auto target = sign ? 0 : kNumRanks; + if (sum == target) return true; + if (is_last_check) { + printf( + "Mooncake elastic barrier timeout, tag: %d, rank: %d, " + "signal-sum: %d, target: %d\n", + kTag, rank_idx, sum, target); + } + return false; + }); +} + +template +__forceinline__ __device__ void gpu_barrier( + const transport::MooncakeGin& gin, const layout::WorkspaceLayout& workspace, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& sm_idx, const int& thread_idx, bool do_scaleout = true, + bool do_scaleup = true) { + if constexpr (kFlushStores) gin.flush(); + if constexpr (kSyncAtStart) { + local_grid_sync(workspace, + thread_idx); + } + + do_scaleout &= kNumScaleoutRanks > 1; + do_scaleup &= kNumScaleupRanks > 1; + if (do_scaleup && !do_scaleout) { + mooncake_barrier_wo_local_sync(gin, workspace, scaleup_rank_idx, + sm_idx, thread_idx); + } else if (do_scaleout && !do_scaleup) { + mooncake_barrier_wo_local_sync( + gin, workspace, scaleout_rank_idx, sm_idx, thread_idx); + } else { + const int global_rank = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + mooncake_barrier_wo_local_sync< + transport::WorldTeam, kNumScaleoutRanks * kNumScaleupRanks, kNumSMs, + kNumThreads, kNumTimeoutCycles, kTag>(gin, workspace, global_rank, + sm_idx, thread_idx); + } + + if constexpr (kSyncAtEnd) { + local_grid_sync(workspace, + thread_idx); + } +} + +} // namespace mooncake::elastic::comm diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh new file mode 100644 index 0000000000..1850361106 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh @@ -0,0 +1,115 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +// Make CLion CUDA indexing work +#ifdef __CLION_IDE__ +#define __CUDA_ARCH__ 900 +#define __CUDACC_RDC__ +#define __CUDACC__ +#endif + +// Remove Torch restrictions +#ifdef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_OPERATORS__ +#endif +#ifdef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#endif +#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_BFLOAT162_OPERATORS__ +#undef __CUDA_NO_BFLOAT162_OPERATORS__ +#endif + +#include +#include +#include + +#if defined(MOONCAKE_EP_USE_MUSA) && defined(__MCC__) && \ + !defined(MOONCAKE_EP_MUSA_LDG_DEFINED) +#define MOONCAKE_EP_MUSA_LDG_DEFINED +template +__device__ __forceinline__ dtype_t __ldg(const dtype_t* ptr) { + return *ptr; +} +#endif + +#ifndef DISABLE_SM90_FEATURES +#include +#elif !defined(MOONCAKE_EP_USE_MUSA) +// Ampere does not support FP8 features +#define __NV_E4M3 0 +#define __NV_E5M2 1 +typedef int __nv_fp8_interpretation_t; +typedef int __nv_fp8x4_e4m3; +typedef uint8_t __nv_fp8_storage_t; +#endif + +// Compatibility: 256 bits LD/ST instructions +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(CUDART_VERSION) and \ + CUDART_VERSION >= 13000 +using longlong4_t = longlong4_32a; +#define make_longlong4_t make_longlong4_32a +#else +struct alignas(32) longlong4_t { + long long x, y, z, w; +}; +__device__ __forceinline__ longlong4_t make_longlong4_t(const long long& x, + const long long& y, + const long long& z, + const long long& w) { + return {x, y, z, w}; +} +#endif + +#ifndef EP_NUM_TOPK_IDX_BITS +#define EP_NUM_TOPK_IDX_BITS 64 +#endif + +namespace mooncake { + +#ifndef DISABLE_SM90_FEATURES +constexpr bool kEnableSM90Features = true; +#else +constexpr bool kEnableSM90Features = false; +#endif + +template +struct int_with_bits; +template <> +struct int_with_bits<8> { + using type = int8_t; +}; +template <> +struct int_with_bits<16> { + using type = int16_t; +}; +template <> +struct int_with_bits<32> { + using type = int32_t; +}; +template <> +struct int_with_bits<64> { + using type = int64_t; +}; + +using topk_idx_t = int_with_bits::type; + +union sf_pack_t { + float fp32; + int ue8m0x4; +}; + +constexpr int kNumTMAAlignedBytes = 16; +constexpr int kNumAlignedSFPacks = 16 / sizeof(sf_pack_t); + +// Some communication channel settings +constexpr int kNumMaxChannels = 1024; + +} // namespace mooncake diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh new file mode 100644 index 0000000000..47cfca8401 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh @@ -0,0 +1,277 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kDoExpand, bool kCachedMode, + // NOTES: this channel concept only applies for scale-out ranks + int kNumSMs, int kNumChannels, int kNumWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kNumHiddenBytes, int kNumSFPacks, + int kNumMaxTokensPerRank, int kNumExperts, int kNumTopk, + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumThreads = kNumWarps * 32, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + bool kDoCreateLinkedList = (kNumScaleoutRanks > 1 and not kCachedMode)> +__global__ void __launch_bounds__(kNumThreads, 1) dispatch_copy_epilogue_impl( + void* buffer, void* workspace, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, void* recv_x, sf_pack_t* recv_sf, + topk_idx_t* recv_topk_idx, float* recv_topk_weights, int* recv_src_metadata, + int* channel_linked_list, int num_recv_tokens, + const int recv_sf_token_stride, const int recv_sf_hidden_stride, + const int scaleout_rank_idx, const int scaleup_rank_idx) { + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + + // For top-k index transformations + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + const auto rank_idx = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + const auto expert_start_idx = kNumExpertsPerRank * rank_idx, + expert_end_idx = kNumExpertsPerRank * (rank_idx + 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + const auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumScaleupRanks, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Will block until the main dispatch kernel has finished and all data are + // visible NOTES: PDL is used, please do not use `__ldg` +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // For no CPU sync case, the number of received tokens should be read from + // the GPU tensor + if (num_recv_tokens == kNumMaxTokensPerRank * kNumRanks) + num_recv_tokens = + psum_num_recv_tokens_per_scaleup_rank[kNumScaleupRanks - 1]; + + // Current rank indices should be maintained + int current_rank_idx = -1, stored_psum_num_recv_tokens; + int current_rank_start = 0, current_rank_end = 0; +#pragma unroll + for (int i = global_warp_idx; i < num_recv_tokens; + i += kNumWarps * kNumSMs) { + // Calculate token index in the buffer + while (i >= current_rank_end) { + current_rank_idx += 1; + EP_DEVICE_ASSERT(current_rank_idx < kNumScaleupRanks); + const auto stored_lane_idx = current_rank_idx % 32; + if (stored_lane_idx == 0 and + current_rank_idx + lane_idx < kNumScaleupRanks) + stored_psum_num_recv_tokens = + psum_num_recv_tokens_per_scaleup_rank[current_rank_idx + + lane_idx]; + current_rank_start = current_rank_end; + current_rank_end = + ptx::exchange(stored_psum_num_recv_tokens, stored_lane_idx); + } + const auto buffer_token = + scaleup_buffer.get_rank_buffer(current_rank_idx) + .get_token_buffer(i - current_rank_start); + + // Wait buffer releases + ptx::tma_store_wait(); + __syncwarp(); + + // Issue TMA loads + // Including all stuffs: data, SF, top-k metadata + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + buffer_token.get_base_ptr(), mbarrier_ptr, + tma_buffer.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + tma_buffer.get_num_bytes()); + } + __syncwarp(); + + // Load target expert indices separately to tolerate TMA load latency + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int dst_expert_idx = -1; + if (lane_idx < kNumTopk) + dst_expert_idx = buffer_token.get_topk_idx_ptr()[lane_idx]; + __syncwarp(); + + // Validate target expert indices and store for non-expand mode + const auto in_range = expert_start_idx <= dst_expert_idx and + dst_expert_idx < expert_end_idx; + const auto master_src_topk_idx = + ptx::get_master_lane_idx(ptx::gather(in_range)); + dst_expert_idx = in_range ? dst_expert_idx - expert_start_idx : -1; + EP_DEVICE_ASSERT(ptx::deduplicate(dst_expert_idx, lane_idx) or + dst_expert_idx == -1); + if (not kDoExpand and lane_idx < kNumTopk) + recv_topk_idx[i * kNumTopk + lane_idx] = + static_cast(dst_expert_idx); + __syncwarp(); + + // Calculate target indices in the tensor + int dst_tensor_idx = -1; + if (not kDoExpand and ptx::elect_one_sync()) { + dst_tensor_idx = i; + } else if (kDoExpand and dst_expert_idx >= 0) { + dst_tensor_idx = + atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1); + } + __syncwarp(); + + // Wait for TMA arrival + if (ptx::elect_one_sync()) + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + __syncwarp(); + + // Maintain linked list + if constexpr (kDoCreateLinkedList) { + if (ptx::elect_one_sync()) + channel_linked_list[tma_buffer.get_linked_list_idx_ptr() + [master_src_topk_idx]] = i; + __syncwarp(); + } + + // Issue TMA stores for data + if (kDoExpand ? (dst_tensor_idx >= 0) : ptx::elect_one_sync()) { + ptx::tma_store_1d( + math::advance_ptr(recv_x, static_cast(dst_tensor_idx) * + kNumHiddenBytes), + tma_buffer.get_hidden_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Store SF + if constexpr (kNumSFPacks > 0) { + constexpr auto kNumFullIters = kNumSFPacks / 32; + const bool do_last_iter = + (kNumSFPacks % 32 != 0) and + (kNumFullIters * 32 + lane_idx < kNumSFPacks); + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + + // Load into registers + const auto smem_src_ptr = tma_buffer.get_sf_ptr(); + sf_pack_t reg_src[kNumFullIters + 1]; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) + reg_src[k] = smem_src_ptr[k * 32 + lane_idx]; + if (do_last_iter) + reg_src[kNumFullIters] = + smem_src_ptr[kNumFullIters * 32 + lane_idx]; + + // Prepare strides + const auto recv_sf_token_stride_i64 = + static_cast(recv_sf_token_stride); + const auto recv_sf_hidden_stride_i64 = + static_cast(recv_sf_hidden_stride); + + // Iterate through all valid indices and store into output buffer + auto mask = kDoExpand ? ptx::gather(dst_tensor_idx >= 0) : 1; + while (mask) { + const int valid_lane_idx = __ffs(mask) - 1; + const auto gmem_dst = math::advance_ptr( + recv_sf, + ptx::exchange(dst_tensor_idx, valid_lane_idx) * + (recv_sf_token_stride_i64 * sizeof(sf_pack_t))); +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) + gmem_dst[(k * 32 + lane_idx) * recv_sf_hidden_stride_i64] = + reg_src[k]; + if (do_last_iter) + gmem_dst[(kNumFullIters * 32 + lane_idx) * + recv_sf_hidden_stride_i64] = + reg_src[kNumFullIters]; + mask ^= 1 << valid_lane_idx; + } + } + + // Store the top-k weights + if (kDoExpand and recv_topk_weights != nullptr and + dst_tensor_idx >= 0) { + recv_topk_weights[dst_tensor_idx] = + tma_buffer.get_topk_weights_ptr()[lane_idx]; + } else if (not kDoExpand and recv_topk_weights != nullptr and + lane_idx < kNumTopk) { + // For backward, weights are optional + recv_topk_weights[i * kNumTopk + lane_idx] = + tma_buffer.get_topk_weights_ptr()[lane_idx]; + } + __syncwarp(); + + // Write source token index + // And: + // - Non-hybrid mode: the source scaleup peer rank index and master + // top-k lane index + // - Hybrid mode: the slot index and master top-k lane index + constexpr int kMetadataStride = 2 + kNumTopk; + if (ptx::elect_one_sync()) { + recv_src_metadata[i * kMetadataStride + 0] = + *tma_buffer.get_src_token_global_idx_ptr(); + if constexpr (kNumScaleoutRanks == 1) { + recv_src_metadata[i * kMetadataStride + 1] = + current_rank_idx * kNumTopk + master_src_topk_idx; + } else { + recv_src_metadata[i * kMetadataStride + 1] = + (i - current_rank_start) * kNumTopk + master_src_topk_idx; + } + } + __syncwarp(); + + // Write reduction source indices + if (kDoExpand and lane_idx < kNumTopk) + recv_src_metadata[i * kMetadataStride + 2 + lane_idx] = + dst_tensor_idx; + __syncwarp(); + } + + // Maintain linked list's ending + // Or you can understand it as writing the tail at once + if constexpr (kDoCreateLinkedList) { + constexpr int kNumScaleupRanksPerLane = + math::constexpr_ceil_div(kNumScaleupRanks, 32); + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + for (int i = global_warp_idx; i < kNumChannels; + i += kNumSMs * kNumWarps) { +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + if (const auto k = j * 32 + lane_idx; + j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) { + channel_linked_list + [*workspace_layout.get_channel_scaleup_tail_ptr(i, k)] = + -1; + + // Clean for combine usages + *workspace_layout.get_channel_scaleup_tail_ptr(i, k) = 0; + } + } + __syncwarp(); + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh new file mode 100644 index 0000000000..323f81d5fb --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh @@ -0,0 +1,171 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include + +namespace mooncake::elastic { + +// Slot preassignment runs in the active scale-up domain. Hybrid scale-out +// forwarding is handled by the hybrid dispatch kernel. +template +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_deterministic_prologue_impl(topk_idx_t* topk_idx, + int* rank_count_buffer, + int* dst_buffer_slot_idx, + const int num_tokens, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumScaleupRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, + "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = sm_idx * kNumWarps + warp_idx; + + // Token region the current warp is responsible for + const auto num_tokens_per_warp = + math::ceil_div(num_tokens, kNumSMs * kNumWarps); + const auto start_token_idx = global_warp_idx * num_tokens_per_warp; + const auto end_token_idx = + min(start_token_idx + num_tokens_per_warp, num_tokens); + + // Group configs + // NOTES: Group refers to the tokens that each warp handles concurrently + constexpr int kNumTokensPerGroup = 32 / kNumTopk; + const auto token_idx_offset = lane_idx / kNumTopk; + const unsigned token_mask = ((1u << kNumTopk) - 1) + << (token_idx_offset * kNumTopk); + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k"); + + // Shared memory for reduction + // NOTES: Each warp owns separate shared memory region for separate sum. + extern __shared__ int8_t smem[]; + const auto rank_count_global_psum = math::advance_ptr(smem, 0); + const auto rank_count_warp_sum = math::advance_ptr( + rank_count_global_psum, + (kNumScaleupRanks + warp_idx * kNumScaleupRanks) * sizeof(int)); + const auto rank_count_warp_psum = math::advance_ptr( + rank_count_warp_sum, kNumWarps * kNumScaleupRanks * sizeof(int)); + + // Initialize to zero before reduce + for (int i = thread_idx; i < kNumScaleupRanks * (1 + 2 * kNumWarps); + i += kNumThreads) + reinterpret_cast(smem)[i] = 0; + __syncthreads(); + + // Util functions + const auto map_expert_to_rank_idx = [&](const int& expert_idx) { + return expert_idx >= 0 ? expert_idx / kNumExpertsPerRank : -1; + }; + const auto is_unique = [&](const int& rank_idx) { + return ((ptx::match(rank_idx) & token_mask) >> lane_idx) == 1; + }; + const auto count_ones_before = [&](const unsigned& mask, + const int& bit_idx) { + return __popc(mask & ((1u << bit_idx) - 1)); + }; + const auto get_other_rank_count_warp_sum = [&](const int& other_warp_idx) { + // NOTES: pass negative num_bytes to advance pointer + return math::advance_ptr( + rank_count_warp_sum, + (other_warp_idx - warp_idx) * kNumScaleupRanks * sizeof(int)); + }; + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = + lane_idx < kNumTopk * kNumTokensPerGroup and + token_idx < end_token_idx; + const int expert_idx = + is_active_thread + ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) + : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_sum[deduped_rank_idx] += __popc(rank_idx_mask); + } + __syncthreads(); + + // Get block sum and store to global + for (int rank_idx = thread_idx; rank_idx < kNumScaleupRanks; + rank_idx += kNumThreads) { + int rank_count_block_sum = 0; + for (int i = 0; i < kNumWarps; i++) + rank_count_block_sum += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_buffer[sm_idx * kNumScaleupRanks + rank_idx] = + rank_count_block_sum; + } + cooperative_groups::this_grid().sync(); + + // Get the prefix sum before the current SM + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = 0; + for (int i = warp_idx; i < sm_idx; i += kNumWarps) + rank_count += rank_count_buffer[i * kNumScaleupRanks + rank_idx]; + atomicAdd_block(rank_count_global_psum + rank_idx, rank_count); + } + __syncthreads(); + + // Get each warp's prefix sum + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = rank_count_global_psum[rank_idx]; + for (int i = 0; i < warp_idx; i++) + rank_count += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_warp_psum[rank_idx] = rank_count; + } + __syncwarp(); + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = + lane_idx < kNumTopk * kNumTokensPerGroup and + token_idx < end_token_idx; + const auto expert_idx = + is_active_thread + ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) + : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Store to target buffer + const auto stored_dst_slot_idx = + deduped_rank_idx >= 0 + ? rank_count_warp_psum[deduped_rank_idx] + + count_ones_before(rank_idx_mask, lane_idx) + : -1; + const auto value = + stored_dst_slot_idx >= 0 + ? scaleup_rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx + : -1; + if (is_active_thread) + dst_buffer_slot_idx[i * kNumTopk + lane_idx] = value; + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_psum[deduped_rank_idx] += __popc(rank_idx_mask); + __syncwarp(); + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh new file mode 100644 index 0000000000..761392dde4 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh @@ -0,0 +1,512 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template > +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, + float* topk_weights, topk_idx_t* copied_topk_idx, + int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + int* dst_buffer_slot_idx, const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + const device::CommCtx comm_ctx, void* buffer, void* workspace, + void* mapped_host_workspace, const int rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + EP_STATIC_ASSERT(kNumExperts % kNumRanks == 0, + "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + + // Workspaces + const auto workspace_layout = + layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + const auto host_workspace_layout = layout::WorkspaceLayout( + mapped_host_workspace, 1, kNumRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared + // memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = + kNumNotifyThreads > 0 ? math::constexpr_align(kNumRanks + kNumExperts, + kNumNotifyThreads) * + sizeof(int) + : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, + "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode 0)>( + sm_idx, warp_idx - kNumNotifyWarps, warp_idx < kNumNotifyWarps); + const auto gin = transport::MooncakeGin(comm_ctx, qp_idx, sharing_mode, + kNumQPs, 0, 0, 0, kNumRanks); + + // Barrier without TMA store flush, without prologue grid sync + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, + // please take care of the `thread_idx` + int *rank_count = rank_expert_count, + *expert_count = rank_expert_count + kNumRanks; +#pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + for (int i = global_warp_idx; i < num_tokens; + i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = + lane_idx < kNumTopk ? static_cast(__ldg( + topk_idx + i * kNumTopk + lane_idx)) + : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = + dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + +// Do full-grid reduction +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add( + workspace_layout.get_notify_reduction_workspace_ptr() + i, + counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { +// Reduce all SM's count +// Wait all SMs' arrival +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while< + kNumTimeoutCycles>(true, [=](const bool& is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout.get_notify_reduction_workspace_ptr() + + i); + if ((status >> 32) == kNumSMs) { + // Write into shared memory + // Write into send buffer if with RDMA + const auto encoded = math::encode_decode_positive( + static_cast(status & 0xffffffffll)); + rank_expert_count[i] = encoded; + if constexpr (not kIsScaleupNVLink) + workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + encoded; + + // Clean for the next usage + workspace_layout + .get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP notify (GPU reduction) timeout, rank: " + "%d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + rank_idx, kNumRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // TODO: for further optimization, we can fuse rank and expert + // counters Issue scaleup rank count writes to peers + for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) { + // Rank counters + const auto dst_rank_counter = + workspace_layout.get_scaleup_rank_count_ptr() + + rank_idx; + gin.put_value(dst_rank_counter, + static_cast(rank_count[i]), i, + 0); + } + __syncwarp(); + + // Issue scaleup expert count writes to peers + if constexpr (kIsScaleupNVLink) { + // NVLink per-element copy + // We don't use TMA as the dtype of shared memory and global is + // different + for (int i = thread_idx; i < kNumExperts; + i += kNumNotifyThreads) { + const auto idx = kNumExpertsPerRank * rank_idx + + (i % kNumExpertsPerRank); + gin.put_value( + workspace_layout.get_scaleup_expert_count_ptr() + + idx, + static_cast(expert_count[i]), + i / kNumExpertsPerRank); + } + } else { + // RDMA bulk copy + for (int i = thread_idx; i < kNumRanks; + i += kNumNotifyThreads) { + const auto src_ptr = + workspace_layout.get_scaleup_expert_count_ptr() + + kNumExpertsPerRank * i; + const auto dst_ptr = + workspace_layout.get_scaleup_expert_count_ptr() + + kNumExpertsPerRank * rank_idx; + gin.put(dst_ptr, src_ptr, + kNumExpertsPerRank * sizeof(int64_t), i); + } + } + + // This is necessary, as the waited results will rewrite the shared + // memory + ptx::named_barrier(kNotifyBarrierIndex); + + // Wait for rank and expert count + const auto start_clock = clock64(); + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while( + [=](const bool& is_last_check) { + // NOTES: the global memory type has 64 bits + const auto count = static_cast< + int>(ptx::ld_volatile( + workspace_layout + .get_scaleup_rank_expert_count_ptr() + + i)); + const auto decoded = + math::encode_decode_positive(count); + if (math::is_decoded_positive_ready(decoded)) { + workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + 0; + rank_expert_count[i] = decoded; + return true; + } + + if (is_last_check) + printf( + "DeepEP notify timeout, rank: %d, thread: %d, " + "count: %d\n", + rank_idx, i, decoded); + return false; + }, + start_clock); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Reduce expert count and add stats + for (int i = thread_idx; i < kNumExpertsPerRank; + i += kNumNotifyThreads) { + int sum = 0; +#pragma unroll + for (int j = 0; j < kNumRanks; ++j) + sum += expert_count[j * kNumExpertsPerRank + i]; + expert_count[i] = math::align(sum, kExpertAlignment); + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr) + atomicAdd(cumulative_local_expert_recv_stats + i, sum); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Write host workspace + if constexpr (kDoCPUSync) { + for (int i = thread_idx; i < kNumRanks + kNumExpertsPerRank; + i += kNumNotifyThreads) { + host_workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + math::encode_decode_positive(rank_expert_count[i]); + } + __syncwarp(); + } + + // Do prefix sum by the warps + // NOTES: we may have fast implementation with `cub::BlockScan`, but + // it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, + const int is_exclusive) { + int psum = 0; +#pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = + (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = + psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, + kNumRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, + kNumExpertsPerRank, 1); + } + } + } else { + const int dispatch_warp_idx = warp_idx - kNumNotifyWarps; + + // Buffer layouts + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout( + token_layout, kNumDispatchWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)) + .get_rank_buffer(dispatch_warp_idx) + .get_token_buffer(0); + auto recv_buffer = layout::BufferLayout( + token_layout, kNumRanks, kNumMaxTokensPerRank, buffer); + auto send_buffer = + layout::BufferLayout(token_layout, 1, kNumMaxTokensPerRank, + recv_buffer.get_buffer_end_ptr()); + recv_buffer = recv_buffer.get_rank_buffer(rank_idx); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Iterate all tokens + const auto token_start = dispatch_warp_idx * kNumSMs + sm_idx; + const auto token_stride = kNumDispatchWarps * kNumSMs; + for (int token_idx = token_start; token_idx < num_tokens; + token_idx += token_stride) { + const auto token_i64_idx = static_cast(token_idx); + + // Wait TMA store arrivals + ptx::tma_store_wait(); + __syncwarp(); + + // Issue data TMA + ptx::tma_load_1d_warp( + tma_buffer.get_hidden_ptr(), + math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes, lane_idx); + __syncwarp(); + + // Issue SF TMA or cp.async + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr( + sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) { + ptx::cp_async_ca( + gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca( + gmem_src_ptr + + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, + "Insufficient lanes for loading top-k indices"); + int stored_dst_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = + __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = + static_cast(uncasted_dst_expert_idx); + stored_dst_rank_idx = dst_expert_idx >= 0 + ? dst_expert_idx / kNumExpertsPerRank + : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = + __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = + uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + // Please ensure no TMA buffer shared memory writes after this part + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = + rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = __ldg( + dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx); + stored_dst_slot_idx = stored_dst_slot_idx >= 0 + ? (stored_dst_slot_idx - + rank_idx * kNumMaxTokensPerRank) + : -1; + } else { + if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and + stored_dst_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd( + workspace_layout.get_scaleup_atomic_sender_counter() + + stored_dst_rank_idx, + 1); + if (lane_idx < kNumTopk) { + const auto value = stored_dst_slot_idx >= 0 + ? rank_idx * kNumMaxTokensPerRank + + stored_dst_slot_idx + : -1; + dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = + value; + } + } + __syncwarp(); + + // Wait TMA load arrival + // NOTES: this arrive must be after the + // `ptx::cp_async_mbarrier_arrive` + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // TMA store to send buffer + auto send_buffer_ptr = + send_buffer.get_token_buffer(token_idx).get_base_ptr(); + if constexpr (not kIsScaleupNVLink) { + if (ptx::elect_one_sync()) + ptx::tma_store_1d(send_buffer_ptr, + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + } + + // Issue TMA NVLink stores + EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid top-k selection"); + const auto dst_ptr = + stored_dst_slot_idx >= 0 + ? gin.get_sym_ptr( + recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + stored_dst_rank_idx) + : nullptr; + if (dst_ptr != nullptr) + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + + // Issue RDMA put + if constexpr (not kIsScaleupNVLink) { + // Wait the send buffer store to arrive + ptx::tma_store_wait<1>(); + __syncwarp(); + + // NOTES: we should skip the NVLink accessible ranks + if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) { + gin.put( + recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + send_buffer_ptr, tma_buffer.get_num_bytes(), + stored_dst_rank_idx); + } + __syncwarp(); + } + } + } + + // Barrier to ensure data arrival + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Trigger the copy epilogue kernel +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Clean atomic counters + EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh new file mode 100644 index 0000000000..26e1c8b984 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh @@ -0,0 +1,81 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include + +#ifndef EP_STATIC_ASSERT +#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) +#endif + +#define EPExceptionWithLineInfo(name, message) \ + EPException(name, __FILE__, __LINE__, message) + +#ifndef EP_HOST_ASSERT +#define EP_HOST_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + throw EPException("Assertion", __FILE__, __LINE__, #cond); \ + } \ + } while (0) +#endif + +#ifndef EP_HOST_UNREACHABLE +#define EP_HOST_UNREACHABLE(reason) \ + (throw EPException("Assertion", __FILE__, __LINE__, reason)) +#endif + +#ifndef EP_DEVICE_ASSERT +#define EP_DEVICE_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, \ + __LINE__, #cond); \ + asm("trap;"); \ + } \ + } while (0) +#endif + +#ifndef EP_UNIFIED_ASSERT +#ifdef __CUDA_ARCH__ +#define EP_UNIFIED_ASSERT(cond) EP_DEVICE_ASSERT(cond) +#else +#define EP_UNIFIED_ASSERT(cond) EP_HOST_ASSERT(cond) +#endif +#endif + +#ifndef CUDA_RUNTIME_CHECK +#define CUDA_RUNTIME_CHECK(cmd) \ + do { \ + const auto e = (cmd); \ + if (e != cudaSuccess) { \ + std::stringstream ss; \ + ss << static_cast(e) << " (" << cudaGetErrorName(e) << ", " \ + << cudaGetErrorString(e) << ")"; \ + throw EPException("CUDA runtime", __FILE__, __LINE__, ss.str()); \ + } \ + } while (0) +#endif + +#ifndef CUDA_DRIVER_CHECK +#define CUDA_DRIVER_CHECK(cmd) \ + do { \ + const auto e = (cmd); \ + if (e != CUDA_SUCCESS) { \ + std::stringstream ss; \ + const char *name, *info; \ + lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \ + ss << static_cast(e) << " (" << name << ", " << info << ")"; \ + throw EPException("CUDA driver", __FILE__, __LINE__, ss.str()); \ + } \ + } while (0) +#endif + +#ifndef NCCL_CHECK +#define NCCL_CHECK(cmd) \ + do { \ + (void)(cmd); \ + } while (0) +#endif diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh new file mode 100644 index 0000000000..1689e8f01f --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh @@ -0,0 +1,787 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs, + int kNumScaleupWarps, int kNumForwardWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kHidden, int kNumMaxTokensPerRank, + int kNumExperts, int kNumTopk, int kNumQPs, int64_t kNumTimeoutCycles, + int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, + 32), + int kNumScaleupUpdateInterval = 3, int kNumChannelsPerSM = kNumForwardWarps, + int kNumChannels = kNumChannelsPerSM * kNumSMs, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumWarps = kNumScaleupWarps + kNumForwardWarps, + int kNumThreads = kNumWarps * 32, + int kNumHiddenBytes = kHidden * sizeof(nv_bfloat16), + bool kUseScaleoutRankLayout = + use_rank_layout(), + bool kUseScaleupRankLayout = + use_rank_layout(), + int kNumTokensInScaleoutLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumScaleoutRanks, kNumTopk>(), + int kNumTokensInScaleupLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumScaleupRanks, kNumTopk>()> +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + int* token_metadata_at_forward, + int* channel_linked_list, + const device::CommCtx comm_ctx, void* buffer, + void* workspace, const int scaleout_rank_idx, + const int scaleup_rank_idx, int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr bool kDoExpandedSend = + not kAllowMultipleReduction and kUseExpandedLayout; + + // Combine vector type selection + using combine_vec_t = typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = kNumHiddenBytes / sizeof(combine_vec_t); + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = + __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumScaleupRanks - 1); + + // Token layouts + const auto token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + + // TMA buffers + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + + // All the buffer layouts + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleupLayout, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleoutLayout, kNumMaxTokensPerRank, + scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_send_buffer = layout::BufferLayout( + token_layout, kAllowMultipleReduction ? 1 : kNumTopk, + kNumChannels * (kNumScaleoutRanks * kNumMaxTokensPerChannel), + scaleout_recv_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-up and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Mooncake Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode( + sm_idx, warp_idx % kNumChannelsPerSM); + const auto gin = transport::MooncakeGin( + comm_ctx, qp_idx, sharing_mode, kNumQPs, scaleout_rank_idx, + scaleup_rank_idx, kNumScaleupRanks, kNumRanks); + + // Global parallel barriers for scale-out subteam and scale-up subteam + // NOTES: this barrier needs a grid sync, as there are channel scale-up tail + // cleaning before + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx); + + // Adjust register count at certain cases + // TODO: support more cases, or try to make channel count more aligned + // DeepEP's register redistribution uses setmaxnreg, which is not accepted + // by ptxas for the SM90 target used by current Mooncake NV validation. + // Keep the official role split but disable this SM100-only optimization. + constexpr bool kAdjustRegisters = false; + constexpr int kNumRegistersForScaleupWarps = 40; + constexpr int kNumRegistersForForwardWarps = + 256 - kNumRegistersForScaleupWarps; + + // Different warp roles + if (warp_idx < kNumScaleupWarps) { + const auto channel_idx = sm_idx * kNumChannelsPerSM + warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_dealloc(); + + // Shift into the right buffer if using rank layout + if constexpr (kUseScaleupRankLayout) + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) + EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Tail issuer + // `st.release.sys` is pretty slow, so do it by an interval + int update_counter = 0; + int stored_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + int stored_old_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, scaleup_rank_idx); + const auto update_tails = [&](const bool& finish = false) { + ++update_counter; + if (finish or update_counter == kNumScaleupUpdateInterval) { + // Wait all TMA stores to finish + ptx::tma_store_wait(); + __syncwarp(); + +// Issue +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + if (const auto j = i * 32 + lane_idx; + i < (kNumScaleupRanksPerLane - 1) or + j < kNumScaleupRanks) { + // NOTES: save some traffic with + // `stored_old_num_tokens_sent` Also, we cannot rewrite + // a finished slot, if the peer is going to clean it + if (stored_num_tokens_sent[i] != + stored_old_num_tokens_sent[i]) + ptx::st_release_sys( + gin.get_sym_ptr( + tail_ptr, j), + stored_num_tokens_sent[i]); + stored_old_num_tokens_sent[i] = + stored_num_tokens_sent[i]; + } + } + update_counter = 0; + } + __syncwarp(); + }; + + // Shape of `channel_linked_list`: `[kNumChannels, + // kNumMaxTokensPerChannel + 1, kNumScaleupRanks]` Iterate until all + // scale-up peers finish + int dst_scaleup_rank_idx = channel_idx; + int stored_ll_idx[kNumScaleupRanksPerLane] = {}, + stored_token_idx[kNumScaleupRanksPerLane] = {}; +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + stored_token_idx[i] = -1; + while (true) { +// Load token indices in the list +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + const auto j = i * 32 + lane_idx; + stored_token_idx[i] = + i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks + ? __ldg( + channel_linked_list + + channel_idx * + (kNumScaleoutRanks * kNumMaxTokensPerChannel + + 1) * + kNumScaleupRanks + + stored_ll_idx[i] * kNumScaleupRanks + j) + : -1; + } + __syncwarp(); + + // Check whether all ranks are finished + bool exited = true; +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + exited &= ptx::all(stored_token_idx[i] < 0); + if (exited) break; + + // Process tokens for all ranks together using bitmask to skip + // inactive ranks + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, + "Too many scale-up ranks for 64-bit mask"); + using mask_t = std::conditional_t<(kNumScaleupRanks <= 32), + uint32_t, uint64_t>; + mask_t wip_mask = 0; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + wip_mask |= + static_cast(ptx::gather(stored_token_idx[j] >= 0)) + << (j * 32); + while (wip_mask) { + // Find next active rank after `dst_scaleup_rank_idx` + // (round-robin) + const auto start = + (dst_scaleup_rank_idx + 1) % kNumScaleupRanks; + const auto hi_mask = (wip_mask >> start) << start; + dst_scaleup_rank_idx = + hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + wip_mask ^= static_cast(1) << dst_scaleup_rank_idx; + + // Exchange token index from the owning lane using static + // partition iteration + int token_idx = -1; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto src_lane_idx = dst_scaleup_rank_idx - j * 32; + token_idx = src_lane_idx == lane_idx ? stored_token_idx[j] + : token_idx; + } + token_idx = ptx::exchange(token_idx, dst_scaleup_rank_idx % 32); + + // Get source metadata and decide the destination buffer + constexpr int kMetadataStride = 2 + kNumTopk; + const auto src_global_token_idx = + __ldg(src_metadata + token_idx * kMetadataStride + 0); + const auto src_token_idx = + src_global_token_idx % kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = + src_global_token_idx / + (kNumMaxTokensPerRank * kNumScaleupRanks); + auto token_buffer = [&]() { + if constexpr (kUseScaleupRankLayout) { + const auto src_slot_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + 1) / + kNumTopk; + return scaleup_buffer.get_token_buffer(src_slot_idx); + } else { + const auto master_topk_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + 1) % + kNumTopk; + return scaleup_buffer.get_rank_buffer(master_topk_idx) + .get_token_buffer(src_scaleout_rank_idx * + kNumMaxTokensPerRank + + src_token_idx); + } + }(); + token_buffer.set_base_ptr( + gin.get_sym_ptr( + token_buffer.get_base_ptr(), dst_scaleup_rank_idx)); + + // Some checks + EP_STATIC_ASSERT( + kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, + "Invalid hidden"); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no-expand, expand + no-reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = + not kUseExpandedLayout or (kAllowMultipleReduction and + __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = token_idx; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = ptx::exchange( + stored_topk_slot_idx, + ptx::get_master_lane_idx(reduce_valid_mask)); + + // Directly load + if (ptx::elect_one_sync()) { + const auto load_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + }); + + // Reduce into shared memory + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * + static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }); + ptx::tma_store_fence(); + __syncwarp(); + } else { +// No local reduction, send all data (expanded send) +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + int topk_slot_idx = + ptx::exchange(stored_topk_slot_idx, k); + if (topk_slot_idx < 0) continue; + + if (ptx::elect_one_sync()) { + // Load + const auto load_ptr = math::advance_ptr( + x, static_cast(kDoExpandedSend + ? topk_slot_idx + : token_idx) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + load_ptr, mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, + phase); + // NOTES: We don't need to care about `topk_weights` + // since we are in expand mode + + // Store + const auto dst_token_buffer = + scaleup_buffer.get_rank_buffer(k) + .get_token_buffer(src_scaleout_rank_idx * + kNumMaxTokensPerRank + + src_token_idx); + ptx::tma_store_1d( + gin.get_sym_ptr( + dst_token_buffer.get_base_ptr(), + dst_scaleup_rank_idx), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + } + } + + // Write top-k weights + if (not kUseExpandedLayout and topk_weights != nullptr and + lane_idx < kNumTopk) { + const float value = + __ldg(topk_weights + (token_idx * kNumTopk + lane_idx)); + tma_buffer.get_topk_weights_ptr()[lane_idx] = value; + ptx::tma_store_fence(); + } + __syncwarp(); + + // Issue TMA stores into remote scale-up buffer + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and ptx::elect_one_sync()) { + // Wait TMA arrival (only for non-reduced cases) + if (no_local_reduce) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + + // Issue stores + ptx::tma_store_1d(token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_num_tokens_sent[j] += + (j * 32 + lane_idx) == dst_scaleup_rank_idx; + __syncwarp(); + } + + // Update the tails together + // NOTES: TMA wait is inside + update_tails(); + +// Move linked list +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + stored_ll_idx[i] += (stored_token_idx[i] >= 0); + } + + // Update for the unissued ones + update_tails(true); + } else { + const auto forward_warp_idx = warp_idx - kNumScaleupWarps; + const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_alloc(); + + // Shift into the right buffer + scaleout_send_buffer = scaleout_send_buffer.get_channel_buffer< + kNumScaleoutRanks * kNumMaxTokensPerChannel>(channel_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, + // kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, + // kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += + channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * + kNumForwardMetadataDims); + + // Overlap TMA stores and reduction + int last_src_scaleout_rank_idx = -1; + int last_is_token_last_in_chunk = 0; + void* last_recv_token_buffer_ptr = nullptr; + void* last_send_token_buffer_ptr = nullptr; + const auto flush_last_tma_and_issue_rdma = [&]() { + if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) { + ptx::tma_store_wait(); + + // Issue only if not local rank + if (last_src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + last_recv_token_buffer_ptr, last_send_token_buffer_ptr, + token_layout.get_num_bytes(), + last_src_scaleout_rank_idx, + last_is_token_last_in_chunk ? 0 : 0); + } + } + __syncwarp(); + }; + + // Replay the dispatch + int stored_num_tokens_recv[kNumScaleupRanksPerLane] = {}, + stored_cached_scaleup_tail[kNumScaleupRanksPerLane] = {}; + for (int i = 0;; ++i) { + const auto src_token_global_idx = + __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims); + const auto is_token_last_in_chunk = __ldg( + token_metadata_at_forward + i * kNumForwardMetadataDims + 1); + const auto src_rank_idx = + src_token_global_idx / kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = src_rank_idx / kNumScaleupRanks; + const auto src_token_idx = + src_token_global_idx % kNumMaxTokensPerRank; + auto stored_src_scaleup_rank_idx = + lane_idx < kNumTopk + ? __ldg(token_metadata_at_forward + + i * kNumForwardMetadataDims + 2 + lane_idx) + : -1; + auto stored_src_slot_idx = lane_idx < kNumTopk + ? __ldg(token_metadata_at_forward + + i * kNumForwardMetadataDims + + 2 + kNumTopk + lane_idx) + : -1; + if (src_token_global_idx < 0) break; + + // Scaleup rank mask + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_mask = + ptx::reduce_or(stored_src_scaleup_rank_idx >= 0 + ? (mask_t(1) << stored_src_scaleup_rank_idx) + : mask_t(0)); + bool stored_is_scaleup_rank_needed[kNumScaleupRanksPerLane]; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_is_scaleup_rank_needed[j] = + (scaleup_mask >> (j * 32 + lane_idx)) & 1; + + // Wait all tails to arrive + comm::timeout_while([&](const bool& + is_last_check) { + bool arrived = true; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + arrived &= not stored_is_scaleup_rank_needed[j] or + stored_num_tokens_recv[j] < + stored_cached_scaleup_tail[j]; + if (ptx::all(arrived)) return true; + +// Reload cached +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto k = j * 32 + lane_idx; + stored_cached_scaleup_tail[j] = + j < (kNumScaleupRanksPerLane - 1) or + k < kNumScaleupRanks + ? ptx::ld_acquire_sys( + workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, k)) + : -1; + } + + // Timeout + if (is_last_check) { +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + printf( + "DeepEP combine (scale-up wait) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, recv: %d, tail: %d " + "(wait=%d)\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, channel_idx, + j * 32 + lane_idx, stored_num_tokens_recv[j], + stored_cached_scaleup_tail[j], + stored_is_scaleup_rank_needed[j]); + } + } + return false; + }); + +// Increase received count +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_num_tokens_recv[j] += + static_cast(stored_is_scaleup_rank_needed[j]); + + if constexpr (not kAllowMultipleReduction) { + // Cases where multiple reduction is disabled. We need to + // forward all data from scaleup peers to scaleout peers + // TODO: Let scale-up warps directly put data into + // `send_buffer`? + const auto src_slot_idx = + src_scaleout_rank_idx * kNumMaxTokensPerRank + + src_token_idx; + auto topk_valid_mask = + kUseExpandedLayout + ? ptx::gather(stored_src_scaleup_rank_idx >= 0) + : ptx::gather( + ptx::deduplicate(stored_src_scaleup_rank_idx, + lane_idx) and + stored_src_scaleup_rank_idx >= + 0); // Deduplicate w.r.t. scaleup rank index + // if expanded mode is disabled + if (ptx::elect_one_sync()) { +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + if ((topk_valid_mask & (1u << k)) == 0u) continue; + + // Issue TMA load, and wait + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + scaleup_buffer.get_rank_buffer(k) + .get_token_buffer(src_slot_idx) + .get_base_ptr(), + mbarrier_ptr, + token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx( + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // Issue TMA store, and wait + const auto recv_buffer_ptr = + scaleout_recv_buffer.get_rank_buffer(k) + .get_token_buffer(src_token_idx) + .get_base_ptr(); + const auto send_buffer_ptr = + src_scaleout_rank_idx == scaleout_rank_idx + ? recv_buffer_ptr + : scaleout_send_buffer.get_rank_buffer(k) + .get_token_buffer(i) + .get_base_ptr(); + ptx::tma_store_1d(send_buffer_ptr, + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue IBGDA + topk_valid_mask ^= 1u << k; + if (src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + recv_buffer_ptr, send_buffer_ptr, + token_layout.get_num_bytes(), + src_scaleout_rank_idx, + topk_valid_mask == 0 and is_token_last_in_chunk + ? 0 + : 0); + } + } + } + __syncwarp(); + } else { + // NOTES: we must do deduplicate and only add once from one rank + auto reduce_valid_mask = ptx::gather( + ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and + stored_src_scaleup_rank_idx >= 0); + + // Calculate the source buffer index + int stored_src_buffer_idx = 0; + if constexpr (kUseScaleupRankLayout) { + stored_src_buffer_idx = + stored_src_scaleup_rank_idx * + scaleup_buffer.num_max_tokens_per_rank + + stored_src_slot_idx; + } else { + const auto src_slot_idx = + src_scaleout_rank_idx * kNumMaxTokensPerRank + + src_token_idx; + stored_src_buffer_idx = + stored_src_slot_idx == -1 + ? -1 + : lane_idx * + scaleup_buffer.num_max_tokens_per_rank + + src_slot_idx; + } + + // Preprocess top-k indices + int topk_slot_idx[kNumTokensInScaleupLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_src_buffer_idx, idx); + }); + + // Do reduce + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return static_cast( + scaleup_buffer.get_token_buffer(slot_idx, true) + .get_base_ptr()); + }, + /* Wait buffer release */ + [=]() { flush_last_tma_and_issue_rdma(); }); + + // Merge topk weights + // NOTES: the slot indices must follow the master lane + stored_src_buffer_idx = ptx::exchange( + stored_src_buffer_idx, ptx::get_master_lane_idx(ptx::match( + stored_src_scaleup_rank_idx))); + if (not kUseExpandedLayout and + stored_src_scaleup_rank_idx >= 0) { + tma_buffer.get_topk_weights_ptr()[lane_idx] = + scaleup_buffer + .get_token_buffer(stored_src_buffer_idx, true) + .get_topk_weights_ptr()[lane_idx]; + } + ptx::tma_store_fence(); + __syncwarp(); // Necessary to let the leader lane see the + // writes + + // Assign send and receive buffers + // NOTES: as we only have 1 destination, we will use "send" as + // "recv" for local transfer + int scaleout_recv_buffer_rank_idx; + if constexpr (kUseScaleoutRankLayout) { + scaleout_recv_buffer_rank_idx = scaleout_rank_idx; + } else { + const int src_topk_idx = ptx::get_master_lane_idx( + ptx::gather(stored_src_scaleup_rank_idx >= 0)); + scaleout_recv_buffer_rank_idx = src_topk_idx; + } + const auto recv_token_buffer = + scaleout_recv_buffer + .get_rank_buffer(scaleout_recv_buffer_rank_idx) + .get_token_buffer(src_token_idx); + const auto send_token_buffer = + src_scaleout_rank_idx == scaleout_rank_idx + ? recv_token_buffer + : scaleout_send_buffer.get_token_buffer(i); + + // Write into scale-out send buffer or local rank recv buffer + // bypass + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Record RDMA info to issue later + last_src_scaleout_rank_idx = src_scaleout_rank_idx; + last_is_token_last_in_chunk = is_token_last_in_chunk; + last_recv_token_buffer_ptr = recv_token_buffer.get_base_ptr(); + last_send_token_buffer_ptr = send_token_buffer.get_base_ptr(); + } + } + + // Issue the last RDMA + if constexpr (kAllowMultipleReduction) flush_last_tma_and_issue_rdma(); + +// Clean scaleup tails +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto k = j * 32 + lane_idx; + if (j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) + *workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k) = + 0; + } + __syncwarp(); + + // Update, wait and clean + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid ranks"); + if (lane_idx < kNumScaleoutRanks) { + // Update remote tails + const auto expected_signal = math::pack2(1, 0); + gin.red_add_rel( + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, scaleout_rank_idx), + expected_signal, lane_idx); + + // Wait tail arrival + const auto wait_ptr = + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx); + comm::timeout_while([=](const bool& + is_last_check) { + const auto signal = ptx::ld_acquire_sys(wait_ptr); + if (signal == expected_signal) { + // Clean for next usages + *wait_ptr = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP combine (scale-out wait all) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, signal: %lld, expected: %lld\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, + kNumScaleupRanks, channel_idx, lane_idx, signal, + expected_signal); + } + return false; + }); + } + __syncwarp(); + } + + // No barrier at epilogue +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh new file mode 100644 index 0000000000..9211911d87 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh @@ -0,0 +1,888 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kDoCPUSync, bool kReuseSlotIndices, int kNumSMs, int kNumNotifyWarps, + int kNumScaleoutWarps, int kNumForwardWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kNumHiddenBytes, int kNumSFPacks, + int kNumMaxTokensPerRank, int kNumExperts, int kNumTopk, + int kExpertAlignment, int kNumQPs, int64_t kNumTimeoutCycles, + int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, + 32), + int kNumChannelsPerSM = kNumScaleoutWarps, + int kNumChannels = kNumScaleoutWarps * kNumSMs, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + int kScaleoutUpdateInterval = 3, + int kNumSlotsPerForwardChunk = kScaleoutUpdateInterval, + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumNotifyThreads = kNumNotifyWarps * 32, + int kNumScaleoutSendThreads = kNumScaleoutWarps * 32, + int kNumForwardThreads = kNumForwardWarps * 32, + int kNumThreads = kNumNotifyThreads + kNumScaleoutSendThreads + + kNumForwardThreads> +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, + float* topk_weights, topk_idx_t* copied_topk_idx, + int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + int* dst_buffer_slot_idx, + int* token_metadata_at_forward, const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + // TODO(NCCL): so many params, plans to optimize? + const device::CommCtx comm_ctx, void* buffer, + void* workspace, void* mapped_host_workspace, + const int scaleout_rank_idx, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, + "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + EP_STATIC_ASSERT(kNumScaleoutWarps == kNumForwardWarps, + "Invalid warp size"); + + // Utils + // NOTES: a warp is a channel (different channels may share QPs) + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto rank_idx = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + const auto host_workspace_layout = + layout::WorkspaceLayout(mapped_host_workspace, kNumScaleoutRanks, + kNumScaleupRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared + // memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = + kNumNotifyThreads > 0 ? math::constexpr_align(kNumRanks + kNumExperts, + kNumNotifyThreads) * + sizeof(int) + : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, + "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // Mooncake Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode 0)>( + sm_idx, (warp_idx - kNumNotifyWarps) % kNumChannelsPerSM, + warp_idx < kNumNotifyWarps); + const auto gin = transport::MooncakeGin( + comm_ctx, qp_idx, sharing_mode, kNumQPs, scaleout_rank_idx, + scaleup_rank_idx, kNumScaleupRanks, kNumRanks); + + // Global parallel barriers for scale-out subteam and scale-up subteam + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx); + + // The golden layout during the whole process for both scale-out and forward + // warps + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout( + token_layout, kNumScaleoutWarps + kNumForwardWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)) + .get_rank_buffer(warp_idx - kNumNotifyWarps) + .get_token_buffer(0); + + // All the buffers + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumScaleupRanks, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + auto scaleout_send_buffer = + layout::BufferLayout(token_layout, 1, kNumMaxTokensPerRank, + scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumScaleoutRanks, kNumChannels * kNumMaxTokensPerChannel, + scaleout_send_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-out and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (warp_idx >= kNumNotifyWarps and ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, + // please take care of the `thread_idx` + int *rank_count = rank_expert_count, + *expert_count = rank_expert_count + kNumRanks; +#pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = sm_idx * kNumNotifyWarps + warp_idx; + for (int i = global_warp_idx; i < num_tokens; + i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = + lane_idx < kNumTopk ? static_cast(__ldg( + topk_idx + i * kNumTopk + lane_idx)) + : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = + dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + +// Do full-grid reduction +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add( + workspace_layout.get_notify_reduction_workspace_ptr() + i, + counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { +// Reduce all SM's count +// Wait all SMs' arrival +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while([=](const bool& + is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout.get_notify_reduction_workspace_ptr() + + i); + if ((status >> 32) == kNumSMs) { + // Encode and write into the send buffer + workspace_layout + .get_scaleout_rank_expert_count_ptr()[i] = + math::encode_decode_positive(status & + 0xffffffffll); + + // Clean for the next usage + workspace_layout + .get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP hybrid notify (GPU reduction) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Issue scaleout writes to peers + EP_STATIC_ASSERT( + kReuseSlotIndices or kNumScaleoutRanks <= kNumNotifyThreads, + "kNumScaleoutRanks must be less than kNumNotifyThreads"); + if (thread_idx < kNumScaleoutRanks) { + const auto dst_scaleout_rank_idx = thread_idx; + gin.put( + workspace_layout.get_scaleout_rank_count_ptr( + scaleout_rank_idx), + workspace_layout.get_scaleout_rank_count_ptr( + dst_scaleout_rank_idx), + kNumScaleupRanks * sizeof(int), dst_scaleout_rank_idx, 0); + gin.put( + workspace_layout.get_scaleout_expert_count_ptr( + scaleout_rank_idx), + workspace_layout.get_scaleout_expert_count_ptr( + dst_scaleout_rank_idx), + kNumExpertsPerScaleout * sizeof(int), + dst_scaleout_rank_idx); + } + __syncwarp(); + + // Util functions to get metadata from scale-out peers + // NOTES: this is correct as RDMA operations has a minimum write + // granularity of 1024 bytes (a whole integer write is atomic) + const auto recv_and_reduce = [=](const auto& get_ptr_func, + const bool& is_expert_reduction = + false) -> int { + int count = 0; +#pragma unroll + for (int j = 0; j < kNumScaleoutRanks; ++j) { + const auto ptr = get_ptr_func(j); + int decoded; + comm::timeout_while< + kNumTimeoutCycles>([&](const bool& is_last_check) { + decoded = math::encode_decode_positive( + ptx::ld_acquire_sys(ptr)); + if (math::is_decoded_positive_ready(decoded)) + return true; + + if (is_last_check) { + printf( + "DeepEP hybrid notify (scale-out %s reduction) " + "timeout, " + "scale-out: %d, scale-up: %d, " + "thread: %d, wait scale-out: %d, decoded: %d\n", + is_expert_reduction ? "expert" : "rank", + scaleout_rank_idx, scaleup_rank_idx, thread_idx, + j, decoded); + } + return false; + }); + + // Add and clean for next usages + count += decoded, *ptr = 0; + } + return count; + }; + +// Write into all scale-up peers' rank-level counters +#pragma unroll + for (int i = thread_idx; i < kNumScaleupRanks; + i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = + recv_and_reduce([=](const int& scaleout_peer_idx) { + return workspace_layout + .get_scaleout_rank_count_ptr( + scaleout_peer_idx, i); + }); + + // Write into the remote scale-up peer + const int64_t counter = + (static_cast(kNumScaleupRanks) << 32ll) | count; + gin.put_value( + workspace_layout.get_scaleup_rank_count_ptr() + + scaleup_rank_idx, + counter, i); + } + __syncwarp(); + +// Atomic add into all scale-up peers' expert-level counters +#pragma unroll + for (int i = thread_idx; i < kNumExpertsPerScaleout; + i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = recv_and_reduce( + [=](const int& scaleout_peer_idx) { + return workspace_layout + .get_scaleout_expert_count_ptr( + scaleout_peer_idx, i); + }, + true); + + // Write into the remote scale-up peer + const int64_t counter = (1ll << 32ll) | count; + const auto dst_scaleup_rank_idx = i / kNumExpertsPerRank; + const auto expert_idx_in_dst_rank = i % kNumExpertsPerRank; + gin.red_add_rel( + workspace_layout.get_scaleup_expert_count_ptr() + + expert_idx_in_dst_rank, + counter, dst_scaleup_rank_idx); + } + // There are shared memory reads above, a barrier is necessary + ptx::named_barrier(kNotifyBarrierIndex); + + // NOTES: from now on, the `rank` and `expert`s size change into the + // local size + expert_count = rank_expert_count + kNumScaleupRanks; + + // Wait local counters to be ready + // NOTES: here we only care the prefix sum by scale-up peers (used + // for later epilogue), not all ranks + EP_STATIC_ASSERT( + kNumNotifyWarps == 0 or kNumScaleupRanks + kNumExpertsPerRank <= + kNumNotifyWarps * 32, + "Insufficient notify threads"); + comm::timeout_while( + thread_idx < kNumScaleupRanks + kNumExpertsPerRank, + [&](const bool& is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout + .get_scaleup_rank_expert_count_ptr() + + thread_idx); + if ((status >> 32ull) == kNumScaleupRanks) { + // Clean GPU workspace and write into host workspace + const auto count = + static_cast(status & 0xffffffffll); + const auto aligned_count = math::align( + count, thread_idx < kNumScaleupRanks + ? 1 + : kExpertAlignment); + + workspace_layout.get_scaleup_rank_expert_count_ptr< + false>()[thread_idx] = 0; + if constexpr (kDoCPUSync) { + host_workspace_layout + .get_scaleup_rank_expert_count_ptr< + false>()[thread_idx] = + math::encode_decode_positive(aligned_count); + } + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr and + thread_idx >= kNumScaleupRanks) + atomicAdd(cumulative_local_expert_recv_stats + + (thread_idx - kNumScaleupRanks), + count); + + // Save for later prefix sum calculation + rank_expert_count[thread_idx] = aligned_count; + return true; + } + + if (is_last_check) { + printf( + "DeepEP hybrid notify (scale-up reduction) timeout," + "scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), + kNumScaleupRanks); + } + return false; + }); + ptx::named_barrier(kNotifyBarrierIndex); + + // Do prefix sum by the warps of the first SM + // NOTES: we may have fast implementation with `cub::BlockScan`, but + // it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, + const int is_exclusive) { + int psum = 0; +#pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = + (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = + psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, + kNumScaleupRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, + kNumExpertsPerRank, 1); + } + } + } else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) { + const int scaleout_warp_idx = warp_idx - kNumNotifyWarps; + const int channel_idx = sm_idx * kNumChannelsPerSM + scaleout_warp_idx; + scaleout_recv_buffer = + scaleout_recv_buffer.get_rank_buffer(scaleout_rank_idx); + scaleout_recv_buffer = + scaleout_recv_buffer.get_channel_buffer( + channel_idx); + + // Channel metadata maintenance + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, + "Invalid number of scale-out ranks"); + int stored_scaleout_tail = 0, stored_old_scaleout_tail = 0; + const auto update_scaleout_tail = [&](const bool& finish_flag = false) { + if (lane_idx < kNumScaleoutRanks and + (stored_scaleout_tail >= + stored_old_scaleout_tail + kScaleoutUpdateInterval or + finish_flag)) { + const auto signaled_tail = math::pack2( + finish_flag, stored_scaleout_tail); + const auto ptr = + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, scaleout_rank_idx); + const auto old_signaled_tail = + math::pack2(0, stored_old_scaleout_tail); + + // NOTES: the "release" scope will be `sys` for the local rank + // (we may involve NVLink so not `gpu`) For RDMA requests, + // "release" is ensured by "atomic" + gin.red_add_rel( + ptr, signaled_tail - old_signaled_tail, lane_idx, + transport::kRedAddReleaseLowWordLast); + stored_old_scaleout_tail = stored_scaleout_tail; + } + __syncwarp(); + }; + + // Preload next token + const auto preload_next_token = [&](const int& token_idx) { + if (token_idx >= num_tokens) return; + + // Issue TMA load + const auto token_i64_idx = static_cast(token_idx); + if (ptx::elect_one_sync()) { + ptx::tma_load_1d( + tma_buffer.get_hidden_ptr(), + math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + + // Issue SF `cp.async` + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr( + sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) { + ptx::cp_async_ca( + gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca( + gmem_src_ptr + + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + }; + + // Iterate all tokens + preload_next_token(channel_idx); + for (int token_idx = channel_idx; token_idx < num_tokens; + token_idx += kNumChannels) { + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, + "Insufficient lanes for loading top-k indices"); + int stored_dst_scaleout_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = + __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = + static_cast(uncasted_dst_expert_idx); + stored_dst_scaleout_rank_idx = + dst_expert_idx >= 0 + ? dst_expert_idx / kNumExpertsPerScaleout + : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = + __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = + uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = + rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + const auto stored_old_slot_idx = ptx::exchange( + stored_scaleout_tail, stored_dst_scaleout_rank_idx >= 0 + ? stored_dst_scaleout_rank_idx + : 0); + if (ptx::deduplicate(stored_dst_scaleout_rank_idx, lane_idx) and + stored_dst_scaleout_rank_idx >= 0) + stored_dst_slot_idx = stored_old_slot_idx; + + // Update scale-out tail + const auto scaleout_rank_mask = + ptx::reduce_or(stored_dst_scaleout_rank_idx >= 0 + ? (1u << stored_dst_scaleout_rank_idx) + : 0u); + stored_scaleout_tail += (scaleout_rank_mask >> lane_idx) & 1; + + // Wait TMA arrival and issue the TMA store into send buffer + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // So if no ranks will go by RDMA, we skip the send buffer + // stores + if (scaleout_rank_mask ^ (1 << scaleout_rank_idx)) { + ptx::tma_store_1d( + scaleout_send_buffer.get_token_buffer(token_idx) + .get_base_ptr(), + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + } + } + __syncwarp(); + + // Local rank can be bypassed + if (stored_dst_slot_idx >= 0 and + stored_dst_scaleout_rank_idx == scaleout_rank_idx) { + ptx::tma_store_1d( + scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + } + ptx::tma_store_commit(); + ptx::tma_store_wait(); + __syncwarp(); + + // Preload the next token (overlapping with the IBGDA issues) + preload_next_token(token_idx + kNumChannels); + + // Issue IBGDA requests + if (stored_dst_slot_idx >= 0 and + stored_dst_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + scaleout_send_buffer.get_token_buffer(token_idx) + .get_base_ptr(), + tma_buffer.get_num_bytes(), + stored_dst_scaleout_rank_idx, 0); + } + __syncwarp(); + + // Issue scale-out tail update + update_scaleout_tail(); + } + + // Flush unflushed tails + update_scaleout_tail(true); + } else { + const int forward_warp_idx = + warp_idx - (kNumNotifyWarps + kNumScaleoutWarps); + const int channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + scaleout_recv_buffer = + scaleout_recv_buffer.get_channel_buffer( + channel_idx); + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, + // kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, + // kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += + channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * + kNumForwardMetadataDims); + + // Shape of `dst_buffer_slot_idx`: `[kNumChannels, kNumScaleoutRanks, + // kNumMaxTokensPerChannel, kNumTopk]` + dst_buffer_slot_idx += + channel_idx * + (kNumScaleoutRanks * kNumMaxTokensPerChannel * kNumTopk); + + // Transform linked list index + const auto transform_linked_list_idx = [=](const int& idx) { + constexpr int kNumTokensInLinkedList = + kNumMaxTokensPerChannel * kNumScaleoutRanks + 1; + return channel_idx * (kNumTokensInLinkedList * kNumScaleupRanks) + + idx * kNumScaleupRanks + scaleup_rank_idx; + }; + + // Forward tokens from scale-out ranks + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Too many scale-out ranks"); + int num_tokens_processed = 0; + int stored_scaleout_old_tail_idx = 0; + int stored_scaleup_send_counters[kNumScaleupRanksPerLane] = {}; + int stored_finish_flag = lane_idx >= kNumScaleoutRanks; + int stored_scaleout_tail_idx = 0; + int recv_scaleout_rank_idx = channel_idx % kNumScaleoutRanks; + uint32_t wip_mask; + while ((wip_mask = ptx::gather(stored_scaleout_tail_idx > + stored_scaleout_old_tail_idx or + stored_finish_flag == 0))) { + // Pick next rank in round-robin + const auto offset = + (recv_scaleout_rank_idx + 1) % kNumScaleoutRanks; + const auto hi_mask = (wip_mask >> offset) << offset; + recv_scaleout_rank_idx = + hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + + // Wait for this rank to have data (or finish) + comm::timeout_while([&](const bool& + is_last_check) { + const uint32_t arrived_or_finished = + stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or + stored_finish_flag > 0; + if (ptx::exchange(arrived_or_finished, recv_scaleout_rank_idx)) + return true; + + // Timeout + if (is_last_check) { + if (lane_idx < kNumScaleoutRanks) { + printf( + "DeepEP hybrid dispatch (forwarding) timeout, " + "scale-out: %d, scale-up: %d, " + "channel: %d, lane: %d, old scale-out tail: %d, " + "scale-out tail: (%d, %d)\n", + scaleout_rank_idx, scaleup_rank_idx, channel_idx, + lane_idx, stored_scaleout_old_tail_idx, + stored_finish_flag, stored_scaleout_tail_idx); + } + return false; + } + + // Read new signaled tails + if (lane_idx < kNumScaleoutRanks) { + const auto signaled_tail = ptx::ld_acquire_sys( + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx)); + math::unpack2(signaled_tail, + stored_finish_flag, + stored_scaleout_tail_idx); + } + __syncwarp(); + return false; + }); + + // Process one chunk from the current rank + const auto start_slot_idx = ptx::exchange( + stored_scaleout_old_tail_idx, recv_scaleout_rank_idx); + const auto end_slot_idx = std::min( + ptx::exchange(stored_scaleout_tail_idx, recv_scaleout_rank_idx), + start_slot_idx + kNumSlotsPerForwardChunk); + if (lane_idx == recv_scaleout_rank_idx) + stored_scaleout_old_tail_idx = end_slot_idx; + + const auto recv_buffer = + scaleout_recv_buffer.get_rank_buffer(recv_scaleout_rank_idx); + for (int slot_idx = start_slot_idx; slot_idx < end_slot_idx; + ++slot_idx) { + const auto token_buffer = + recv_buffer.get_token_buffer(slot_idx); + + // Wait TMA arrival + ptx::tma_store_wait(); + __syncwarp(); + + // TMA load into shared memory + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + token_buffer.get_base_ptr(), mbarrier_ptr, + token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx( + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // Read top-k indices + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int stored_dst_scaleup_rank_idx = -1; + auto dst_expert_idx = + lane_idx < kNumTopk + ? tma_buffer.get_topk_idx_ptr()[lane_idx] + : -1; + dst_expert_idx -= scaleout_rank_idx * kNumExpertsPerScaleout; + stored_dst_scaleup_rank_idx = + 0 <= dst_expert_idx and + dst_expert_idx < kNumExpertsPerScaleout + ? dst_expert_idx / kNumExpertsPerRank + : -1; + + // Write the per-scaleup channel index for this token + int linked_list_idx = -1; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto src_lane_idx = + stored_dst_scaleup_rank_idx - j * 32; + const bool valid = 0 <= src_lane_idx and src_lane_idx < 32; + const auto exchanged = + ptx::exchange(stored_scaleup_send_counters[j], + valid ? src_lane_idx : 0); + linked_list_idx = valid ? exchanged : linked_list_idx; + } + if (not kReuseSlotIndices and lane_idx < kNumTopk) { + tma_buffer.get_linked_list_idx_ptr()[lane_idx] = + transform_linked_list_idx(linked_list_idx); + ptx::tma_store_fence(); + } + __syncwarp(); + + // Deduplicate for scale-up ranks + int stored_dst_slot_idx = -1; + const auto dst_slot_idx_ptr = + dst_buffer_slot_idx + + recv_scaleout_rank_idx * + (kNumMaxTokensPerChannel * kNumTopk) + + slot_idx * kNumTopk; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = + __ldg(dst_slot_idx_ptr + lane_idx); + } else { + // Deduplicate for NVLink ranks + if (ptx::deduplicate(stored_dst_scaleup_rank_idx, + lane_idx) and + stored_dst_scaleup_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd( + workspace_layout + .get_scaleup_atomic_sender_counter() + + stored_dst_scaleup_rank_idx, + 1); + } + __syncwarp(); + + // Issue TMAs + if (stored_dst_slot_idx >= 0) { + const auto dst_ptr = + gin.get_sym_ptr( + scaleup_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + stored_dst_scaleup_rank_idx); + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Add per-scale-up counter + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, + "Invalid number of scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_send_mask = ptx::reduce_or( + stored_dst_scaleup_rank_idx >= 0 + ? (mask_t(1) << stored_dst_scaleup_rank_idx) + : mask_t(0)); +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_scaleup_send_counters[j] += + (scaleup_send_mask >> (j * 32 + lane_idx)) & 1; + + // Record metadata at forward + if constexpr (not kReuseSlotIndices) { + EP_STATIC_ASSERT(kNumTopk <= 32, + "Invalid number of selections"); + const auto metadata_ptr = + token_metadata_at_forward + + num_tokens_processed * kNumForwardMetadataDims; + + // Source token index and last token index flag + if (ptx::elect_one_sync()) { + metadata_ptr[0] = + tma_buffer.get_src_token_global_idx_ptr()[0]; + metadata_ptr[1] = slot_idx == (end_slot_idx - 1); + } + + // Second, original top-k indices and destination slots + if (lane_idx < kNumTopk) { + metadata_ptr[2 + lane_idx] = + stored_dst_scaleup_rank_idx; + metadata_ptr[2 + kNumTopk + lane_idx] = + stored_dst_slot_idx; + dst_slot_idx_ptr[lane_idx] = stored_dst_slot_idx; + } + } + num_tokens_processed += 1; + __syncwarp(); + } + } + + // Assign the source token index part of the metadata into `-1` as an + // ending mark + if (not kReuseSlotIndices and ptx::elect_one_sync()) + token_metadata_at_forward[num_tokens_processed * + kNumForwardMetadataDims] = -1; + __syncwarp(); + + // Update linked list's ending position + if constexpr (not kReuseSlotIndices) { + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, scaleup_rank_idx); +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + if (const auto j = i * 32 + lane_idx; + i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) { + ptx::st_relaxed_sys( + gin.get_sym_ptr(tail_ptr, j), + transform_linked_list_idx( + stored_scaleup_send_counters[i])); + } + } + } + __syncwarp(); + + // Clean tails for next usages + if (lane_idx < kNumScaleoutRanks) + *workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx) = 0; + __syncwarp(); + } + + // Scale-up barrier to ensure data arrival + // As scale-out tokens have already been consumed by forwarders, no need to + // do scale-out barrier again + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx, /* do not scale-out */ false, true); + + // Trigger the copy epilogue kernel +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Clean scale-up counters + // All scale-out counters should be cleaned before + EP_STATIC_ASSERT(kNumScaleupRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumScaleupRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh new file mode 100644 index 0000000000..e21c86a170 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh @@ -0,0 +1,78 @@ +#pragma once + +#include + +#include +#include + +namespace mooncake { + +struct ElasticLaunchContext { + void* gdr_buffer = nullptr; + const int32_t* nvlink_available = nullptr; + void* const* ipc_peer_ptrs = nullptr; + void* raddrs = nullptr; + void* rkeys = nullptr; + void* qp_devctxs = nullptr; + const void* rdma_send_signal_buffer = nullptr; + const void* rdma_recv_signal_buffer = nullptr; + void* buffer = nullptr; + void* workspace = nullptr; + void* mapped_host_workspace = nullptr; + int rank = 0; + int num_ranks = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + int num_scaleout_ranks = 1; + int num_scaleup_ranks = 1; + bool is_scaleup_nvlink = true; + int num_qps = 1; + int64_t timeout_cycles = -1; +}; + +void launch_elastic_dispatch_deterministic_prologue( + const int64_t* topk_idx, int* rank_count_buffer, int* dst_buffer_slot_idx, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int scaleup_rank_idx, int num_scaleup_ranks, int num_sms, + int num_smem_bytes, cudaStream_t stream); + +void launch_mooncake_elastic_dispatch( + void* x, void* sf, int64_t* topk_idx, float* topk_weights, + int64_t* copied_topk_idx, int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int* token_metadata_at_forward, int num_tokens, int num_max_tokens_per_rank, + int hidden, int elem_size, int num_sf_packs, int sf_token_stride, + int sf_hidden_stride, int num_experts, int num_topk, int expert_alignment, + int num_sms, int num_channels_per_sm, int num_smem_bytes, bool cached_mode, + bool deterministic, bool do_cpu_sync, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +void launch_mooncake_elastic_dispatch_copy_epilogue( + void* recv_x, void* recv_sf, int64_t* recv_topk_idx, + float* recv_topk_weights, int* recv_src_metadata, int* channel_linked_list, + int num_recv_tokens, int num_max_tokens_per_rank, int hidden, int elem_size, + int num_sf_packs, int recv_sf_token_stride, int recv_sf_hidden_stride, + int num_experts, int num_topk, int num_sms, int num_smem_bytes, + int num_channels, bool do_expand, bool cached_mode, + const ElasticLaunchContext& ctx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, cudaStream_t stream); + +void* launch_mooncake_elastic_combine( + void* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, int* token_metadata_at_forward, + int* channel_linked_list, int num_reduced_tokens, + int num_max_tokens_per_rank, int hidden, int num_experts, int num_topk, + int num_sms, int num_smem_bytes, int num_channels, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +void launch_mooncake_elastic_combine_reduce_epilogue( + void* combined_x, float* combined_topk_weights, int64_t* combined_topk_idx, + int num_combined_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, void* reduce_buffer, void* bias_0, + void* bias_1, int num_sms, int num_smem_bytes, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +} // namespace mooncake diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh new file mode 100644 index 0000000000..f03d5943bb --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh @@ -0,0 +1,374 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include +#include + +namespace mooncake::elastic::layout { + +struct WorkspaceLayout { + void* workspace; + + int num_ranks; + int num_scaleout_ranks, num_scaleup_ranks; + int num_experts, num_experts_per_rank; + + // We want to fix the layout position for all settings, + // so that one buffer can be reused for all cases + static constexpr int kNumMaxRanks = 1024; + static constexpr int kNumMaxExperts = 2048; + static constexpr int kNumMaxExpertsPerRank = 256; + static constexpr int kNumMaxInflightAGRS = 32; + + // Mooncake Device API does not rely on NCCL GIN remote RED on a single + // symmetric signal word. Use per-source-rank signal slots for both phases: + // each sender atomically updates its own slot with release semantics and + // receivers poll the full slot vector. Keep an independent counter/slot + // vector for each logical barrier tag, as hybrid kernels mix world and + // scale-up-only barriers in the same workspace and therefore must not share + // phase/sign state across tags. + static constexpr int kNumBarrierTags = 16; + static constexpr int64_t kNumBarrierBytesPerTag = + sizeof(unsigned long long) + 2 * kNumMaxRanks * sizeof(int); + static constexpr int64_t kNumBarrierSignalBytes = + kNumBarrierTags * kNumBarrierBytesPerTag; + + __forceinline__ __device__ __host__ + WorkspaceLayout(void* workspace, const int& num_scaleout_ranks, + const int& num_scaleup_ranks, const int& num_experts) + : workspace(workspace), + num_ranks(num_scaleout_ranks * num_scaleup_ranks), + num_scaleout_ranks(num_scaleout_ranks), + num_scaleup_ranks(num_scaleup_ranks), + num_experts(num_experts) { + num_experts_per_rank = num_experts / num_ranks; + EP_UNIFIED_ASSERT(num_experts % num_ranks == 0); + EP_UNIFIED_ASSERT(num_ranks <= kNumMaxRanks); + EP_UNIFIED_ASSERT(num_experts <= kNumMaxExperts); + EP_UNIFIED_ASSERT(num_experts_per_rank <= kNumMaxExpertsPerRank); + } + + static int64_t get_num_bytes() { + // Pure NVLink scaleup barrier signals + int64_t num_bytes = 0; + num_bytes += kNumBarrierSignalBytes; + + // Notify reduction workspace + num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t); + + // Scaleup notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int64_t) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int64_t) * 2; + + // Scaleup atomic sender count + num_bytes += kNumMaxRanks * sizeof(int); + + // Scaleout notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int) * 2; + + // Scaleout channel metadata (finish flag and tails) + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t); + + // Channel aggregated into the scaleup domains + // Also reused for channel scaleup tail + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int); + + // Rank send/recv count, for PP prev/next ranks + num_bytes += 2 * 2 * sizeof(int64_t); + + // AGRS signals + num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int); + + // Ensure LDG.256 work + return math::align(num_bytes, 32); + } + + __forceinline__ __device__ __host__ unsigned long long* + get_nvl_barrier_counter_ptr(int tag = 0) const { + EP_UNIFIED_ASSERT(tag >= 0 && tag < kNumBarrierTags); + return math::advance_ptr( + workspace, tag * kNumBarrierBytesPerTag); + } + + __forceinline__ __device__ __host__ int* get_nvl_barrier_signal_ptr( + int tag, int phase) const { + EP_UNIFIED_ASSERT(tag >= 0 && tag < kNumBarrierTags); + EP_UNIFIED_ASSERT(phase >= 0 && phase < 2); + return math::advance_ptr(workspace, + tag * kNumBarrierBytesPerTag + + sizeof(unsigned long long) + + phase * kNumMaxRanks * sizeof(int)); + } + + __forceinline__ __device__ __host__ int64_t* + get_notify_reduction_workspace_ptr() const { + return math::advance_ptr(workspace, kNumBarrierSignalBytes); + } + + template + __forceinline__ __device__ __host__ int64_t* + get_scaleup_rank_expert_count_ptr() const { + const auto base_ptr = math::advance_ptr( + get_notify_reduction_workspace_ptr(), + (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_rank_count_ptr() + const { + return get_scaleup_rank_expert_count_ptr(); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_expert_count_ptr() + const { + return get_scaleup_rank_expert_count_ptr() + + num_scaleup_ranks; + } + + __forceinline__ __device__ __host__ int* get_scaleup_atomic_sender_counter() + const { + return math::advance_ptr( + get_scaleup_rank_expert_count_ptr(), + 2 * (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + } + + template + __forceinline__ __device__ __host__ int* + get_scaleout_rank_expert_count_ptr() const { + const auto base_ptr = math::advance_ptr( + get_scaleup_atomic_sender_counter(), kNumMaxRanks * sizeof(int)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_rank_count_ptr( + const int& scaleout_rank_idx = 0, + const int& scaleup_rank_idx = 0) const { + const auto base_ptr = + get_scaleout_rank_expert_count_ptr(); + return base_ptr + scaleout_rank_idx * num_scaleup_ranks + + scaleup_rank_idx; + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_expert_count_ptr( + const int& scaleout_rank_idx = 0, const int& expert_idx = 0) const { + const auto base_ptr = + get_scaleout_rank_expert_count_ptr() + num_ranks; + return base_ptr + + scaleout_rank_idx * (num_scaleup_ranks * num_experts_per_rank) + + expert_idx; + } + + __forceinline__ __device__ __host__ int64_t* + get_scaleout_channel_signaled_tail_ptr(const int& channel_idx, + const int& scaleout_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_rank_expert_count_ptr(), + (kNumMaxRanks + kNumMaxExperts) * sizeof(int) * 2); + return base_ptr + + (channel_idx * num_scaleout_ranks + scaleout_rank_idx); + } + + __forceinline__ __device__ __host__ int* get_channel_scaleup_tail_ptr( + const int& channel_idx, const int& scaleup_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_channel_signaled_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int64_t)); + return base_ptr + (channel_idx * num_scaleup_ranks + scaleup_rank_idx); + } + + __forceinline__ __device__ __host__ int64_t* get_pp_send_count_ptr( + const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_channel_scaleup_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int64_t* get_pp_recv_count_ptr( + const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_pp_send_count_ptr(0), 2 * sizeof(int64_t)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int* get_agrs_recv_signal_ptr( + const int& slot, const int& rank_idx) const { + const auto base_ptr = math::advance_ptr(get_pp_recv_count_ptr(0), + 2 * sizeof(int64_t)); + return base_ptr + slot * kNumMaxRanks + rank_idx; + } + + __forceinline__ __device__ __host__ int* get_agrs_session_signal_ptr( + const int& rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_agrs_recv_signal_ptr(0, 0), + kNumMaxInflightAGRS * kNumMaxRanks * sizeof(int)); + return base_ptr + rank_idx; + } +}; + +struct TokenLayout { + int num_hidden_bytes, num_sf_bytes; + // NOTES: the top-k index is always 32-bit + bool with_metadata; + int num_topk, num_metadata_bytes; + void* base; + + __forceinline__ __device__ __host__ TokenLayout(const int& num_hidden_bytes, + const int& num_sf_bytes, + const int& num_topk, + const bool& with_metadata, + void* base = nullptr) + : num_hidden_bytes(num_hidden_bytes), + num_sf_bytes(num_sf_bytes), + // Metadata includes: top-k indices, weight and source rank/token + // index + with_metadata(with_metadata), + num_topk(num_topk), + num_metadata_bytes( + num_topk * (sizeof(int) + sizeof(float)) + + (with_metadata ? (1 + num_topk) * sizeof(int) : 0)), + base(base) { + EP_STATIC_ASSERT(sizeof(int) == sizeof(float), + "Invalid size assumption"); + EP_UNIFIED_ASSERT(num_hidden_bytes % ptx::kNumTMAAlignBytes == 0); + } + + template + __forceinline__ __device__ __host__ dtype_t get_num_bytes() const { + const auto num_bytes = + math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) + + math::align(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, + ptx::kNumTMAAlignBytes); + return static_cast(num_bytes); + } + + __forceinline__ __device__ __host__ void* get_base_ptr() const { + return base; + } + + __forceinline__ __device__ __host__ void set_base_ptr(void* ptr) { + base = ptr; + } + + __forceinline__ __device__ __host__ void* get_hidden_ptr() const { + return get_base_ptr(); + } + + __forceinline__ __device__ __host__ sf_pack_t* get_sf_ptr() const { + return math::advance_ptr( + base, math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_metadata_ptr() const { + return math::advance_ptr( + get_sf_ptr(), math::align(num_sf_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_topk_idx_ptr() const { + return get_metadata_ptr(); + } + + __forceinline__ __device__ __host__ float* get_topk_weights_ptr() const { + return math::advance_ptr(get_metadata_ptr(), + num_topk * sizeof(int)); + } + + __forceinline__ __device__ __host__ int* get_src_token_global_idx_ptr() + const { + return math::advance_ptr(get_topk_weights_ptr(), + num_topk * sizeof(float)); + } + + __forceinline__ __device__ __host__ int* get_linked_list_idx_ptr() const { + return get_src_token_global_idx_ptr() + 1; + } + + __forceinline__ __device__ ptx::mbarrier* get_mbarrier_ptr() const { + return math::advance_ptr( + get_metadata_ptr(), + math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes)); + } +}; + +template +struct BufferLayout { + TokenLayout token_layout; + int num_ranks; + int num_max_tokens_per_rank; + + void* base; + + __forceinline__ __device__ __host__ + BufferLayout(const TokenLayout& token_layout, const int& num_ranks, + const int& max_num_tokens_per_rank, void* base = nullptr) + : token_layout(token_layout), + num_ranks(num_ranks), + num_max_tokens_per_rank(max_num_tokens_per_rank), + base(base) {} + + __forceinline__ __device__ __host__ int64_t + get_num_bytes_per_token() const { + return token_layout.get_num_bytes(); + } + + __forceinline__ __device__ __host__ int64_t get_num_bytes_per_rank() const { + return num_max_tokens_per_rank * get_num_bytes_per_token(); + } + + __forceinline__ __device__ __host__ int64_t get_num_bytes() const { + return get_num_bytes_per_rank() * num_ranks; + } + + __forceinline__ __device__ __host__ void* get_buffer_end_ptr() const { + return math::advance_ptr(base, get_num_bytes()); + } + + __forceinline__ __device__ __host__ BufferLayout + get_rank_buffer(const int& rank_idx) const { + return BufferLayout( + token_layout, 1, num_max_tokens_per_rank, + static_cast(base) + get_num_bytes_per_rank() * rank_idx); + } + + template + __forceinline__ __device__ __host__ BufferLayout + get_channel_buffer(const int& channel_idx) const { + EP_UNIFIED_ASSERT(kNumTokensPerChannel > 0); + return BufferLayout( + token_layout, + // Do not use `num_max_tokens_per_rank / kNumTokensPerChannel` as + // the false stride + num_ranks, num_max_tokens_per_rank, + static_cast(base) + + get_num_bytes_per_token() * kNumTokensPerChannel * channel_idx); + } + + __forceinline__ __device__ __host__ TokenLayout + get_token_buffer(const int& token_idx, const bool& global = false) const { + EP_UNIFIED_ASSERT(num_ranks == 1 or global); + return TokenLayout( + token_layout.num_hidden_bytes, token_layout.num_sf_bytes, + token_layout.num_topk, token_layout.with_metadata, + static_cast(base) + + token_layout.get_num_bytes() * + token_idx); + } +}; + +} // namespace mooncake::elastic::layout diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh new file mode 100644 index 0000000000..89dbcd0fcc --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh @@ -0,0 +1,83 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +namespace mooncake::elastic::math { + +template +__forceinline__ __device__ __host__ T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ T align(T a, T b) { + return (kDoCeilAlignment ? ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_align(T a, T b) { + return (kDoCeilAlignment ? constexpr_ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ bool is_decoded_positive_ready( + const dtype_t& value) { + return value >= 0; +} + +template +__forceinline__ __device__ __host__ dtype_t +encode_decode_positive(const dtype_t& value) { + return -value - static_cast(1); +} + +template +__forceinline__ __device__ __host__ dtype_t* advance_ptr( + void* ptr, const int64_t num_bytes) { + return reinterpret_cast(static_cast(ptr) + num_bytes); +} + +__forceinline__ __device__ __host__ ptrdiff_t ptr_diff(const void* ptr, + const void* base) { + return static_cast(ptr) - static_cast(base); +} + +template +__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, + const dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + dtype_b_t packed; + auto unpacked_ptr = reinterpret_cast(&packed); + unpacked_ptr[0] = x, unpacked_ptr[1] = y; + return packed; +} + +template +__device__ __forceinline__ std::tuple unpack2( + const dtype_b_t& packed) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + dtype_a_t x = unpacked_ptr[0], y = unpacked_ptr[1]; + return {x, y}; +} + +template +__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, + dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + x = unpacked_ptr[0], y = unpacked_ptr[1]; +} + +} // namespace mooncake::elastic::math diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh new file mode 100644 index 0000000000..f1fa20650e --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh @@ -0,0 +1,735 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include + +#include +#include + +namespace mooncake::elastic::ptx { + +// Host-side placeholder with the same size/alignment as +// cuda::barrier (a single uint64_t atomic), so that +// sizeof(mbarrier) is consistent across host and device. +struct alignas(8) mbarrier { + uint64_t __placeholder; +}; +using arrival_phase = uint32_t; + +// More than TMA, `longlong4` requires 32 bytes aligned +static constexpr int kNumTMAAlignBytes = 32; + +#ifdef __CUDACC__ + +/// Exceptions +__forceinline__ __device__ void trap() { +#ifdef MOONCAKE_EP_USE_MUSA + return; +#else + asm volatile("trap;"); +#endif +} + +/// Thread layout +__forceinline__ __device__ int get_warp_idx() { + return __shfl_sync(0xffffffff, threadIdx.x / 32, 0); +} + +__forceinline__ __device__ int get_lane_idx() { +#ifdef MOONCAKE_EP_USE_MUSA + return static_cast(threadIdx.x) & 31; +#else + int lane_idx; + asm volatile("mov.s32 %0, %laneid;" : "=r"(lane_idx)); + return lane_idx; +#endif +} + +/// Election +__forceinline__ __device__ int elect_one_sync() { +#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(DISABLE_SM90_FEATURES) && \ + defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + int pred = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %1;\n" + "@%%px mov.s32 %0, 1;\n" + "}\n" + : "+r"(pred) + : "r"(0xffffffff)); + return pred; +#else + return get_lane_idx() == 0; +#endif +} + +/// TMA and `cp.async` +__forceinline__ __device__ void mbarrier_init_with_fence( + mbarrier* ptr, const int& arrive_count = 1) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" ::"r"(arrive_count), + "r"(static_cast(__cvta_generic_to_shared(ptr)))); + asm volatile("fence.mbarrier_init.release.cluster;" ::); +#endif +} + +__forceinline__ __device__ void mbarrier_invalidate(mbarrier* ptr) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_arrive(mbarrier* ptr) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_arrive_and_set_tx( + mbarrier* ptr, const int& num_bytes) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" ::"r"( + num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_wait_and_flip_phase( + mbarrier* ptr, arrival_phase& phase) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + "}" ::"r"(static_cast(__cvta_generic_to_shared(ptr))), + "r"(phase), "r"(0x989680)); +#endif + phase ^= 1; +} + +__forceinline__ __device__ void tma_store_fence() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.async.shared::cta;"); +#endif +} + +template +__forceinline__ __device__ void tma_store_wait() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(kNumRemainingWaits) + : "memory"); +#endif +} + +enum TMACacheHint : int64_t { + kEvictFirst = 0x12f0000000000000ll, + kEvictNormal = 0x1000000000000000ll +}; + +__forceinline__ __device__ void tma_load_1d( + const void* dst_ptr, const void* src_ptr, mbarrier* ptr, + const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictFirst) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + // NOTES: normally, the loaded part will be evicted soon + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::" + "cache_hint [%0], [%1], %2, [%3], %4;\n" ::"r"( + static_cast(__cvta_generic_to_shared(dst_ptr))), + "l"(src_ptr), "r"(num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr))), "l"(hint) + : "memory"); +#else + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_vecs; ++i) dst[i] = src[i]; + } else { + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); + for (int i = 0; i < num_bytes; ++i) dst[i] = src[i]; + } +#endif +} + +__forceinline__ __device__ void tma_load_1d_warp( + const void* dst_ptr, const void* src_ptr, mbarrier* ptr, + const int& num_bytes, const int& lane_idx, + const TMACacheHint& hint = TMACacheHint::kEvictFirst) { +#ifdef MOONCAKE_EP_USE_MUSA + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = lane_idx; i < num_vecs; i += 32) dst[i] = src[i]; + } else { + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); + for (int i = lane_idx; i < num_bytes; i += 32) dst[i] = src[i]; + } +#else + if (elect_one_sync()) tma_load_1d(dst_ptr, src_ptr, ptr, num_bytes, hint); +#endif +} + +__forceinline__ __device__ void tma_store_1d( + const void* dst_ptr, const void* src_ptr, const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictNormal) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + // NOTES: normally, the stored part will be used soon + asm volatile( + "cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], " + "[%1], %2, %3;\n" ::"l"(dst_ptr), + "r"(static_cast(__cvta_generic_to_shared(src_ptr))), + "r"(num_bytes), "l"(hint) + : "memory"); +#else + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_vecs; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile int* src_words = + reinterpret_cast(src + i); + volatile int* dst_words = reinterpret_cast(dst + i); + dst_words[0] = src_words[0]; + dst_words[1] = src_words[1]; + dst_words[2] = src_words[2]; + dst_words[3] = src_words[3]; +#else + dst[i] = src[i]; +#endif + } + } else { +#ifdef MOONCAKE_EP_USE_MUSA + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); +#else + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); +#endif + for (int i = 0; i < num_bytes; ++i) dst[i] = src[i]; + } +#endif +} + +__forceinline__ __device__ void tma_store_commit() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("cp.async.bulk.commit_group;"); +#endif +} + +template +__forceinline__ __device__ void cp_async_ca(const dtype_t* gmem_src, + const dtype_t* smem_dst) { + EP_STATIC_ASSERT( + sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8 or sizeof(dtype_t) == 16, + "Invalid dtype bytes"); +#ifdef MOONCAKE_EP_USE_MUSA + *const_cast(smem_dst) = *gmem_src; +#else + asm volatile( + "cp.async.ca.shared::cta.global.L2::128B [%0], [%1], %2;\n" ::"r"( + static_cast(__cvta_generic_to_shared(smem_dst))), + "l"(gmem_src), "n"(sizeof(dtype_t))); +#endif +} + +__forceinline__ __device__ void cp_async_mbarrier_arrive(mbarrier* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)ptr; +#else + asm volatile("cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +/// Barriers +template +__forceinline__ __device__ void named_barrier(const int& idx) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)idx; + __threadfence_block(); + __syncwarp(); +#else + // Equivalent to `barrier.sync.aligned`, which requires all threads run the + // same location of code + asm volatile("bar.sync %0, %1;" ::"r"(idx), "r"(kNumThreads)); +#endif +} + +/// LD/ST instructions +__forceinline__ __device__ int4 +ldg_with_gez_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value >= 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +__forceinline__ __device__ int4 +ldg_with_gtz_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value > 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.gt.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +__forceinline__ __device__ int4 +ld_with_gez_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value >= 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) +__forceinline__ __device__ longlong4_t +ldg_with_gez_pred(const longlong4_t* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + longlong4_t ret = make_longlong4_t(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s64 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+l"(ret.x), "+l"(ret.y), "+l"(ret.z), "+l"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +__forceinline__ __device__ longlong4_t ldg(const longlong4_t* ptr) { + longlong4_t ret; + asm volatile( + "ld.L1::no_allocate.global.nc.v4.s64 {%0, %1, %2, %3}, [%4];\n\t" + : "=l"(ret.x), "=l"(ret.y), "=l"(ret.z), "=l"(ret.w) + : "l"(ptr) + : "memory"); + return ret; +} +#endif + +__forceinline__ __device__ int4 ldg(const int4* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile int* words = reinterpret_cast(ptr); + int4 ret; + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + return ret; +#else + return __ldg(ptr); +#endif +} + +__forceinline__ __device__ void st_na(int4* ptr, const int4& value) { +#ifdef MOONCAKE_EP_USE_MUSA + volatile int* words = reinterpret_cast(ptr); + words[0] = value.x; + words[1] = value.y; + words[2] = value.z; + words[3] = value.w; +#else + *ptr = value; +#endif +} + +template +__forceinline__ __device__ void st_with_gez_pred(dtype_t* ptr, dtype_t value, + const int& condition) { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4, "Invalid data type"); + auto view = *reinterpret_cast(&value); +#ifdef MOONCAKE_EP_USE_MUSA + if (condition >= 0) *ptr = value; +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %2, 0;\n\t" + " @p st.global.s32 [%0], %1;\n\t" + "}" ::"l"(ptr), + "r"(view), "r"(condition) + : "memory"); +#endif +} + +template +__forceinline__ __device__ dtype_t ld_volatile(const void* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile dtype_t* typed = + reinterpret_cast(ptr); + return *typed; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.volatile.global.u64 %0, [%1];" + : "=l"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +__forceinline__ __device__ void red_add(const int64_t* ptr, + const int64_t& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast( + reinterpret_cast(ptr)), + static_cast(value)); + __threadfence_system(); +#else + // TODO(NVCC): why don't NVCC support `s64`? + // Mooncake elastic consumers can poll these counters from a different CTA + // immediately after the RED producer issues the update. Keep the update as + // a single 64-bit RED so packed counters stay atomic, but use release scope + // instead of a relaxed GPU-scope RED to avoid occasional visibility stalls + // in the notify reduction path. + asm volatile("red.release.gpu.global.add.u64 [%0], %1;" ::"l"(ptr), + "l"(value) + : "memory"); +#endif +} + +__forceinline__ __device__ void red_add_rel_sys(const int* ptr, + const int& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast(ptr), value); + __threadfence_system(); +#else + asm volatile("red.release.sys.global.add.s32 [%0], %1;" ::"l"(ptr), + "r"(value)); +#endif +} + +__forceinline__ __device__ void red_add_rel_sys(const int64_t* ptr, + const int64_t& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast( + reinterpret_cast(ptr)), + static_cast(value)); + __threadfence_system(); +#else + asm volatile("red.release.sys.global.add.u64 [%0], %1;" ::"l"(ptr), + "l"(value)); +#endif +} + +template +__forceinline__ __device__ dtype_t ld_acquire_sys(const dtype_t* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile dtype_t* typed = + reinterpret_cast(ptr); + const dtype_t value = *typed; + __threadfence_system(); + return value; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u32 %0, [%1];" + : "=r"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u64 %0, [%1];" + : "=l"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +template +__forceinline__ __device__ void st_relaxed_sys(void* ptr, dtype_t value) { +#ifdef MOONCAKE_EP_USE_MUSA + *static_cast(ptr) = value; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(ptr), + "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u64 [%0], %1;" ::"l"(ptr), + "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +template +__forceinline__ __device__ void st_release_sys(void* ptr, dtype_t value) { +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); + *static_cast(ptr) = value; + __threadfence_system(); +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(ptr), + "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(ptr), + "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +// Adjust registers +template +__device__ __forceinline__ void warpgroup_reg_alloc() { +#ifndef MOONCAKE_EP_USE_MUSA + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +#endif +} + +template +__device__ __forceinline__ void warpgroup_reg_dealloc() { +#ifndef MOONCAKE_EP_USE_MUSA + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +#endif +} + +/// General fences +__device__ __forceinline__ void fence_acq_rel_sys() { +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); +#else + asm volatile("fence.acq_rel.sys;" ::: "memory"); +#endif +} + +/// Intrinsics +template +__device__ __forceinline__ dtype_t exchange(dtype_t ptr, + const int& src_lane_idx) { + EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, ""); + const auto send_int_values = reinterpret_cast(&ptr); + dtype_t recv_dtype; + auto recv_int_values = reinterpret_cast(&recv_dtype); +#pragma unroll + for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i) + recv_int_values[i] = + __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx); + return recv_dtype; +} + +__device__ __forceinline__ unsigned gather(const bool& value) { + return __ballot_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool all(const bool& value) { + return __all_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool any(const bool& value) { + return __any_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned reduce_or(const unsigned& value) { + return __reduce_or_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned long long reduce_or( + const unsigned long long& value) { + const auto low = __reduce_or_sync(0xffffffff, static_cast(value)); + const auto high = + __reduce_or_sync(0xffffffff, static_cast(value >> 32)); + return (static_cast(high) << 32) | low; +} + +__device__ __forceinline__ int reduce_add(const int& value) { + return __reduce_add_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned match(const int& value) { + return __match_any_sync(0xffffffff, value); +} + +__device__ __forceinline__ int fns(const unsigned& value, const int& offset) { + return __fns(value, 0, offset); +} + +template +__device__ __forceinline__ auto ffs(const dtype_t& value) { + if constexpr (sizeof(dtype_t) == 4) { + return __ffs(static_cast(value)) - 1; + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 8, "Invalid data type"); + return __ffsll(static_cast(value)) - 1; + } +} + +__device__ __forceinline__ int get_master_lane_idx(const unsigned& mask) { +#ifdef MOONCAKE_EP_USE_MUSA + return 31 - __clz(mask); +#else + // Equivalent to `31 - __clz(mask)` + int highest_idx; + asm volatile("bfind.u32 %0, %1;" : "=r"(highest_idx) : "r"(mask)); + return highest_idx; +#endif +} + +__device__ __forceinline__ bool deduplicate(const int& value, + const int& lane_idx) { + return get_master_lane_idx(match(value)) == lane_idx; +} + +__device__ __forceinline__ int warp_inclusive_sum(int value, + const int& lane_idx) { +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const auto synced = __shfl_up_sync(0xffffffff, value, offset); + if (lane_idx >= offset) value += synced; + } + return value; +} + +__device__ __forceinline__ float2 fadd2(const float2& a, const float2& b) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) + return __fadd2_rn(a, b); +#else + return {a.x + b.x, a.y + b.y}; +#endif +} + +__device__ __forceinline__ void accumulate(float2& a, nv_bfloat162 b) { +#ifdef MOONCAKE_EP_USE_MUSA + a.x += __low2float(b); + a.y += __high2float(b); +#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Use `add.rn.f32.bf16` instruction to perform fused (cast + add) operation + // on SM100 + asm("add.rn.f32.bf16 %0, %1, %0;\n" + : "+f"(a.x) + : "h"(*reinterpret_cast(&b.x))); + asm("add.rn.f32.bf16 %0, %1, %0;\n" + : "+f"(a.y) + : "h"(*reinterpret_cast(&b.y))); +#else + const auto [x, y] = __bfloat1622float2(b); + a.x += x, a.y += y; +#endif +} + +#endif + +} // namespace mooncake::elastic::ptx diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh new file mode 100644 index 0000000000..7855ce3a63 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh @@ -0,0 +1,261 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace mooncake::elastic::transport { + +struct WorldTeam {}; +struct ScaleupTeam {}; +struct ScaleoutTeam {}; + +constexpr int kRedAddReleaseHighWordLast = 0; +constexpr int kRedAddReleaseLowWordLast = 1 << 0; + +// Mooncake Device API adapter for DeepEP's NCCL GIN usage. +// +// DeepEP elastic kernels express all remote communication through a small GIN +// surface: symmetric-pointer translation, put, put_value, RED/add style signals +// and QP flushes. Mooncake maps that surface onto Device API semantics: +// +// get_sym_ptr -> mc_route_put, returning local/P2P peer VA or nullptr +// put -> local/P2P warp copy, otherwise mc_rdma_put +// put_value -> local/P2P release store, otherwise mc_rdma_put/mc_signal +// flush -> no-op; Device API operations are ordered by release/fence +// and +// explicit kernel barriers in the imported elastic kernels +// +// Team tags are kept as types so official DeepEP template code can remain close +// to the source while the actual routing is decided by Mooncake CommCtx. +struct MooncakeGin { + device::CommCtx ctx; + int qp_idx = 0; + int sharing_mode = 0; + int qps_per_rank = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + int num_scaleup_ranks = 0; + + __device__ __forceinline__ MooncakeGin( + const device::CommCtx& ctx, int qp_idx, int sharing_mode, int num_qps, + int scaleout_rank_idx = 0, int scaleup_rank_idx = 0, + int num_scaleup_ranks = 0, int num_ranks = 1) + : ctx(ctx), + qp_idx(qp_idx), + sharing_mode(sharing_mode), + qps_per_rank(max(1, num_qps / max(1, num_ranks))), + scaleout_rank_idx(scaleout_rank_idx), + scaleup_rank_idx(scaleup_rank_idx), + num_scaleup_ranks(num_scaleup_ranks) {} + + template + __device__ __forceinline__ int world_rank(int dst_rank) const { + if (num_scaleup_ranks <= 0) return dst_rank; + if constexpr (std::is_same_v) { + return scaleout_rank_idx * num_scaleup_ranks + dst_rank; + } else if constexpr (std::is_same_v) { + return dst_rank * num_scaleup_ranks + scaleup_rank_idx; + } else { + return dst_rank; + } + } + + template + __device__ __forceinline__ bool is_nvlink_accessible(int dst_rank) const { + dst_rank = world_rank(dst_rank); + return dst_rank == ctx.rank || + device::mc_comm_p2p_available(ctx, dst_rank); + } + + template + __device__ __forceinline__ void* get_sym_ptr(void* ptr, + int dst_rank) const { + dst_rank = world_rank(dst_rank); + return device::mc_route_put(ctx, dst_rank, ptr); + } + + template + __device__ __forceinline__ const void* get_sym_ptr(const void* ptr, + int dst_rank) const { + dst_rank = world_rank(dst_rank); + return device::mc_route_put(ctx, dst_rank, const_cast(ptr)); + } + + template + __device__ __forceinline__ void put(void* dst_ptr, const void* src_ptr, + int num_bytes, int dst_rank, + int /*flags*/ = 0) const { + dst_rank = world_rank(dst_rank); + auto routed = device::mc_route_put(ctx, dst_rank, dst_ptr); + if (routed != nullptr) { + const auto src_addr = reinterpret_cast(src_ptr); + const auto dst_addr = reinterpret_cast(routed); + if (((src_addr | dst_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + const auto* src = reinterpret_cast(src_ptr); + auto* dst = reinterpret_cast(routed); + const int num_int4 = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_int4; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + ptx::st_na(dst + i, device::mc_ld_nc(src + i)); +#else + dst[i] = device::mc_ld_nc(src + i); +#endif + } + } else { + auto* dst_bytes = reinterpret_cast(routed); + const auto* src_bytes = + reinterpret_cast(src_ptr); + for (int i = 0; i < num_bytes; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + reinterpret_cast(dst_bytes)[i] = + reinterpret_cast(src_bytes)[i]; +#else + dst_bytes[i] = src_bytes[i]; +#endif + } + } + // `put` is used both by full data-moving warps and by individual + // notify lanes. Do not place a full-warp barrier inside the + // transport primitive: divergent notify calls would deadlock. Each + // participating lane copies the complete payload for its request, + // so a system fence is sufficient to publish the writes. + __threadfence_system(); + } else { + device::mc_rdma_put(ctx, qp_idx, dst_rank, qps_per_rank, src_ptr, + dst_ptr, static_cast(num_bytes), 0); + } + } + + template + __device__ __forceinline__ void put_value(value_t* dst_ptr, value_t value, + int dst_rank, + int flags = 0) const { + dst_rank = world_rank(dst_rank); + auto* routed = + static_cast(device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + device::mc_st_release(reinterpret_cast(routed), + static_cast(value)); + } else { + *routed = value; + __threadfence_system(); + } + } else { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + device::mc_signal(ctx, dst_rank, qp_idx, qps_per_rank, + reinterpret_cast(dst_ptr), + static_cast(value)); + } else { + // Device RDMA WRITE sources must be registered GDR addresses; + // a by-value scalar lives in thread-local storage and is not a + // valid IBGDA source. Current elastic uses remote int64 + // put_value only for single-writer, zeroed notify slots, so a + // split 32-bit RED add is equivalent to writing the packed + // word. + auto* words = reinterpret_cast(dst_ptr); + const auto signed_value = static_cast(value); + const auto low = static_cast( + static_cast(signed_value) & 0xffffffffull); + const auto high = static_cast(signed_value >> 32); + if ((flags & kRedAddReleaseLowWordLast) == 0) { + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + } else { + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + } + } + } + } + + template + __device__ __forceinline__ void red_add_rel(value_t* dst_ptr, value_t value, + int dst_rank, + int flags = 0) const { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + dst_rank = world_rank(dst_rank); + auto* routed = + static_cast(device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + device::mc_atomic_add_release(routed, static_cast(value)); + } else { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + reinterpret_cast(dst_ptr), + static_cast(value)); + } + } else if constexpr (sizeof(value_t) == sizeof(uint64_t) || + sizeof(value_t) == sizeof(int64_t)) { + dst_rank = world_rank(dst_rank); + auto* routed = static_cast( + device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + // Some official elastic paths use the high 32 bits as the + // readiness word (notify counters), while others use the low 32 + // bits as the terminal flag (hybrid channel tails). Splitting + // a 64-bit RED into two 32-bit atomics can therefore publish + // the wrong half first for one of the protocols. Use one + // system- scope 64-bit RED on the routed local/P2P VA so the + // packed value is updated atomically with release ordering. + ptx::red_add_rel_sys(routed, static_cast(value)); + } else { + // Mooncake's current Device API only exposes 32-bit remote + // reduction. Do not emulate the 64-bit add with an RDMA WRITE + // from a thread-local scalar: IBGDA WQEs use the registered GDR + // buffer lkey, so a stack/local address is not a valid DMA + // source on true cross-node runs. Split the packed signal into + // two 32-bit remote reductions instead, publishing the + // readiness word last. Most notify counters use high word as + // the ready count; hybrid channel tails use low word as the + // finish flag. + auto* words = reinterpret_cast(dst_ptr); + const auto signed_value = static_cast(value); + const auto low = static_cast( + static_cast(signed_value) & 0xffffffffull); + const auto high = static_cast(signed_value >> 32); + if ((flags & kRedAddReleaseLowWordLast) == 0) { + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + } else { + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + } + } + } else { + put_value(dst_ptr, value, dst_rank, flags); + } + } + + __device__ __forceinline__ void flush() const { __threadfence_system(); } +}; + +} // namespace mooncake::elastic::transport diff --git a/mooncake-ep/include/mooncake_ep_buffer.h b/mooncake-ep/include/mooncake_ep_buffer.h index 03873c5ca2..9307936c80 100644 --- a/mooncake-ep/include/mooncake_ep_buffer.h +++ b/mooncake-ep/include/mooncake_ep_buffer.h @@ -16,6 +16,7 @@ namespace mooncake { class TransferEngine; +class MooncakeElasticBuffer; // MAX_QP_COUNT is defined in mooncake_ep_configs.cuh (shared with kernel code). @@ -64,6 +65,8 @@ struct BufferPair { struct MooncakeEpBuffer { private: + friend class MooncakeElasticBuffer; + // Device info and communication int device_id; int rank, num_ranks; diff --git a/mooncake-ep/include/mooncake_ep_device.h b/mooncake-ep/include/mooncake_ep_device.h index 9322aea49b..e88417ff7f 100644 --- a/mooncake-ep/include/mooncake_ep_device.h +++ b/mooncake-ep/include/mooncake_ep_device.h @@ -20,8 +20,13 @@ __device__ __forceinline__ ep_fp8x2_storage_t ep_cvt_float2_to_fp8x2(float2 x) { #endif // -- Device intrinsics (MUSA doesn't have __ldg / __activemask) -------------- -#ifndef __ldg -#define __ldg(ptr) (*(ptr)) +#if (defined(__CUDACC__) || defined(__MCC__)) && \ + !defined(MOONCAKE_EP_MUSA_LDG_DEFINED) +#define MOONCAKE_EP_MUSA_LDG_DEFINED +template +__device__ __forceinline__ dtype_t __ldg(const dtype_t* ptr) { + return *ptr; +} #endif #ifndef __activemask #define __activemask() (0xffffffff) diff --git a/mooncake-ep/include/mooncake_ep_exception.cuh b/mooncake-ep/include/mooncake_ep_exception.cuh index 097421f3a2..64061dcd9e 100644 --- a/mooncake-ep/include/mooncake_ep_exception.cuh +++ b/mooncake-ep/include/mooncake_ep_exception.cuh @@ -7,6 +7,8 @@ #define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) #endif +#ifndef MOONCAKE_EP_EXCEPTION_CLASS_DEFINED +#define MOONCAKE_EP_EXCEPTION_CLASS_DEFINED class EPException : public std::exception { private: std::string message = {}; @@ -20,6 +22,7 @@ class EPException : public std::exception { const char* what() const noexcept override { return message.c_str(); } }; +#endif #ifndef CUDA_CHECK #define CUDA_CHECK(cmd) \ diff --git a/mooncake-ep/setup.py b/mooncake-ep/setup.py index 955af759bd..79a1625167 100644 --- a/mooncake-ep/setup.py +++ b/mooncake-ep/setup.py @@ -62,7 +62,10 @@ def existing_dirs(*paths): if use_musa: cuda_libraries = [] - musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"] + musa_defines = [ + "-DUSE_MUSA", + "-DMOONCAKE_EP_USE_MUSA=1", + ] cxx_args += musa_defines # torchada maps the "nvcc" key to "mcc". device_args = [ @@ -113,7 +116,9 @@ def existing_dirs(*paths): sources=[ "src/ep_py.cpp", "src/mooncake_ep_buffer.cpp", + "src/mooncake_ep_elastic_buffer.cpp", "src/mooncake_ep_kernel.cu", + "src/mooncake_ep_elastic_kernel.cu", ], extra_compile_args={"cxx": cxx_args, "nvcc": device_args}, libraries=cuda_libraries, diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index 574ab514c0..a102f5011e 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -1,4 +1,4 @@ -add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu) +add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_elastic_buffer.cpp mooncake_ep_kernel.cu mooncake_ep_elastic_kernel.cu) set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON) target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5) diff --git a/mooncake-ep/src/ep_py.cpp b/mooncake-ep/src/ep_py.cpp index 361df2f96f..02307c0caf 100644 --- a/mooncake-ep/src/ep_py.cpp +++ b/mooncake-ep/src/ep_py.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -13,12 +14,57 @@ namespace mooncake { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("get_ep_buffer_size_hint", &get_ep_buffer_size_hint); + m.def("calculate_elastic_buffer_size", + &MooncakeElasticBuffer::calculate_buffer_size); py::class_(m, "EventHandle") .def(py::init<>()) .def("current_stream_wait", &EventHandle::current_stream_wait) .def("synchronize", &EventHandle::synchronize); + py::class_(m, "ElasticNativeHandle") + .def(py::init<>()) + .def_readwrite("do_expand", &ElasticNativeHandle::do_expand) + .def_readwrite("num_experts", &ElasticNativeHandle::num_experts) + .def_readwrite("expert_alignment", + &ElasticNativeHandle::expert_alignment) + .def_readwrite("num_max_tokens_per_rank", + &ElasticNativeHandle::num_max_tokens_per_rank) + .def_readwrite("num_sms", &ElasticNativeHandle::num_sms) + .def_readwrite("topk_idx", &ElasticNativeHandle::topk_idx) + .def_readwrite( + "psum_num_recv_tokens_per_scaleup_rank", + &ElasticNativeHandle::psum_num_recv_tokens_per_scaleup_rank) + .def_readwrite("psum_num_recv_tokens_per_expert", + &ElasticNativeHandle::psum_num_recv_tokens_per_expert) + .def_readwrite("recv_src_metadata", + &ElasticNativeHandle::recv_src_metadata) + .def_readwrite("recv_layout_range", + &ElasticNativeHandle::recv_layout_range) + .def_readwrite("dst_buffer_slot_idx", + &ElasticNativeHandle::dst_buffer_slot_idx) + .def_readwrite("token_metadata_at_forward", + &ElasticNativeHandle::token_metadata_at_forward) + .def_readwrite("channel_linked_list", + &ElasticNativeHandle::channel_linked_list) + .def_readwrite("num_recv_tokens_per_expert_list", + &ElasticNativeHandle::num_recv_tokens_per_expert_list); + + py::class_(m, "ElasticDispatchOutput") + .def_readonly("recv_x", &ElasticDispatchOutput::recv_x) + .def_readonly("recv_x_scales", &ElasticDispatchOutput::recv_x_scales) + .def_readonly("recv_topk_idx", &ElasticDispatchOutput::recv_topk_idx) + .def_readonly("recv_topk_weights", + &ElasticDispatchOutput::recv_topk_weights) + .def_readonly("handle", &ElasticDispatchOutput::handle) + .def_readonly("event", &ElasticDispatchOutput::event); + + py::class_(m, "ElasticCombineOutput") + .def_readonly("combined_x", &ElasticCombineOutput::combined_x) + .def_readonly("combined_topk_weights", + &ElasticCombineOutput::combined_topk_weights) + .def_readonly("event", &ElasticCombineOutput::event); + m.attr("MAX_QP_COUNT") = pybind11::int_(MAX_QP_COUNT); py::class_(m, "Buffer") @@ -39,6 +85,50 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("combine", &MooncakeEpBuffer::combine) .def("get_next_combine_buffer", &MooncakeEpBuffer::get_next_combine_buffer); + + py::class_(m, "ElasticBuffer") + .def(py::init(), + py::arg("rank"), py::arg("num_ranks"), py::arg("num_buffer_bytes"), + py::arg("num_max_tokens_per_rank"), py::arg("hidden"), + py::arg("num_topk"), py::arg("use_fp8_dispatch"), + py::arg("deterministic"), py::arg("allow_hybrid_mode"), + py::arg("allow_multiple_reduction"), + py::arg("prefer_overlap_with_compute"), py::arg("sl_idx"), + py::arg("num_allocated_qps"), py::arg("num_cpu_timeout_secs"), + py::arg("num_gpu_timeout_secs")) + .def_static("calculate_buffer_size", + &MooncakeElasticBuffer::calculate_buffer_size) + .def("get_physical_domain_size", + &MooncakeElasticBuffer::get_physical_domain_size) + .def("get_logical_domain_size", + &MooncakeElasticBuffer::get_logical_domain_size) + .def("get_theoretical_num_sms", + &MooncakeElasticBuffer::get_theoretical_num_sms) + .def("ibgda_disabled", &MooncakeElasticBuffer::ibgda_disabled) + .def("use_fast_path", &MooncakeElasticBuffer::use_fast_path) + .def("update_local_qpns", &MooncakeElasticBuffer::update_local_qpns) + .def("is_roce", &MooncakeElasticBuffer::is_roce) + .def("sync_ibgda_peers", &MooncakeElasticBuffer::sync_ibgda_peers) + .def("get_mr_info", &MooncakeElasticBuffer::get_mr_info) + .def("get_gid", &MooncakeElasticBuffer::get_gid) + .def("get_local_qpns", &MooncakeElasticBuffer::get_local_qpns) + .def("get_local_lids", &MooncakeElasticBuffer::get_local_lids) + .def("get_ipc_handle", &MooncakeElasticBuffer::get_ipc_handle) + .def("sync_nvlink_ipc_handles", + &MooncakeElasticBuffer::sync_nvlink_ipc_handles) + .def("dispatch", &MooncakeElasticBuffer::dispatch, py::arg("x"), + py::arg("sf"), py::arg("topk_idx"), py::arg("topk_weights"), + py::arg("active_ranks"), py::arg("num_experts"), + py::arg("num_max_tokens_per_rank"), py::arg("expert_alignment"), + py::arg("num_sms"), py::arg("do_expand"), py::arg("do_cpu_sync"), + py::arg("async_with_compute_stream"), + py::arg("cached_handle") = std::nullopt) + .def("combine", &MooncakeElasticBuffer::combine, py::arg("x"), + py::arg("handle"), py::arg("topk_weights"), + py::arg("active_ranks"), py::arg("num_sms"), + py::arg("async_with_compute_stream"), + py::arg("out") = std::nullopt); } } // namespace mooncake diff --git a/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp b/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp new file mode 100644 index 0000000000..7d3d0052cf --- /dev/null +++ b/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp @@ -0,0 +1,632 @@ +#include +#include + +#include +#include +#include +#include + +#include + +namespace mooncake { +namespace { + +int64_t ceil_div_i64(int64_t x, int64_t y) { return (x + y - 1) / y; } + +constexpr int kElasticHybridChannelsPerSm = 4; + +int64_t align_i64(int64_t x, int64_t alignment) { + return ceil_div_i64(x, alignment) * alignment; +} + +int getenv_int(const char* name, int default_value) { + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') return default_value; + return std::max(1, std::atoi(value)); +} + +int hybrid_num_channels(int num_sms) { + return std::max(1, num_sms) * kElasticHybridChannelsPerSm; +} + +int hybrid_num_max_tokens_per_channel(int num_max_tokens_per_rank, + int num_sms) { + return static_cast( + ceil_div_i64(num_max_tokens_per_rank, hybrid_num_channels(num_sms))); +} + +int64_t elastic_workspace_num_bytes() { + constexpr int64_t kNumMaxRanks = 1024; + constexpr int64_t kNumMaxExperts = 2048; + constexpr int64_t kNumMaxChannels = 8 * 160; + constexpr int64_t kNumMaxInflightAGRS = 32; + constexpr int64_t kNumBarrierTags = 16; + + int64_t num_bytes = 0; + num_bytes += kNumBarrierTags * + (sizeof(unsigned long long) + 2 * kNumMaxRanks * sizeof(int)); + num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t); + num_bytes += kNumMaxRanks * sizeof(int64_t) * 2; + num_bytes += kNumMaxExperts * sizeof(int64_t) * 2; + num_bytes += kNumMaxRanks * sizeof(int); + num_bytes += kNumMaxRanks * sizeof(int) * 2; + num_bytes += kNumMaxExperts * sizeof(int) * 2; + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t); + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int); + num_bytes += 2 * 2 * sizeof(int64_t); + num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int); + return align_i64(num_bytes, 32); +} + +int64_t elastic_atomic_scratch_num_bytes() { + return elastic_workspace_num_bytes(); +} + +int device_smem_bytes() { +#ifdef MOONCAKE_EP_USE_MUSA + return 0; +#else + int device = 0; + cudaGetDevice(&device); + int value = 0; + cudaDeviceGetAttribute(&value, cudaDevAttrMaxSharedMemoryPerBlockOptin, + device); + return value > 0 ? value : 98304; +#endif +} + +} // namespace + +ElasticLaunchContext MooncakeElasticBuffer::make_launch_context( + MooncakeEpBuffer& buffer, const ElasticTopology& topology, + void* mapped_host_workspace, int64_t timeout_cycles) { + ElasticLaunchContext ctx; + auto* rdma = buffer.rdma_transport_; + auto* gdr_base = static_cast(buffer.gdr_buffer); + // Mooncake P2P/RDMA Device API translates remote pointers as offsets from + // the registered GDR buffer base. DeepEP elastic writes both `buffer` and + // `workspace` pointers to peer ranks through GIN, so both regions must live + // inside the same peer-visible registered allocation. The elastic buffer + // size reserves `elastic_workspace_num_bytes()` first; use that prefix as + // the workspace. RDMA atomics also need a separate local response area: + // mlx5 atomics write the fetched old value to the WQE local address, so + // reusing the remote signal workspace as `local_atomic_base` can corrupt + // the barrier/signal slots. Reserve an equal-sized scratch prefix after + // the workspace, then place the communication buffer after both prefixes. + const auto workspace_bytes = elastic_workspace_num_bytes(); + const auto atomic_scratch_bytes = elastic_atomic_scratch_num_bytes(); + ctx.gdr_buffer = gdr_base; + ctx.nvlink_available = buffer.p2p_transport_->availableTablePtr(); + ctx.ipc_peer_ptrs = buffer.p2p_transport_->peerPtrsTablePtr(); + ctx.raddrs = rdma ? rdma->raddrsPtr() : nullptr; + ctx.rkeys = rdma ? rdma->rkeysPtr() : nullptr; + ctx.qp_devctxs = rdma ? rdma->qpDevCtxsPtr() : nullptr; + ctx.rdma_send_signal_buffer = gdr_base + workspace_bytes; + ctx.rdma_recv_signal_buffer = gdr_base; + ctx.workspace = gdr_base; + ctx.buffer = gdr_base + workspace_bytes + atomic_scratch_bytes; + ctx.mapped_host_workspace = mapped_host_workspace; + ctx.rank = topology.rank_idx; + ctx.num_ranks = topology.num_ranks; + ctx.scaleout_rank_idx = topology.scaleout_rank_idx; + ctx.scaleup_rank_idx = topology.scaleup_rank_idx; + ctx.num_scaleout_ranks = topology.num_scaleout_ranks; + ctx.num_scaleup_ranks = topology.num_scaleup_ranks; + ctx.is_scaleup_nvlink = true; + ctx.num_qps = buffer.USE_QP_COUNT; + ctx.timeout_cycles = timeout_cycles; + return ctx; +} + +MooncakeElasticBuffer::MooncakeElasticBuffer( + int rank, int num_ranks, int64_t num_buffer_bytes, + int64_t num_max_tokens_per_rank, int64_t hidden, int64_t num_topk, + bool use_fp8_dispatch, bool deterministic, bool allow_hybrid_mode, + bool allow_multiple_reduction, bool prefer_overlap_with_compute, int sl_idx, + int num_allocated_qps, int num_cpu_timeout_secs, int num_gpu_timeout_secs) { + config_.num_max_tokens_per_rank = num_max_tokens_per_rank; + config_.hidden = hidden; + config_.num_topk = num_topk; + config_.use_fp8_dispatch = use_fp8_dispatch; + config_.deterministic = deterministic; + config_.allow_hybrid_mode = allow_hybrid_mode; + config_.allow_multiple_reduction = allow_multiple_reduction; + config_.prefer_overlap_with_compute = prefer_overlap_with_compute; + config_.sl_idx = sl_idx; + config_.num_allocated_qps = num_allocated_qps; + config_.num_cpu_timeout_secs = num_cpu_timeout_secs; + config_.num_gpu_timeout_secs = num_gpu_timeout_secs; + + topology_ = discover_topology(rank, num_ranks, allow_hybrid_mode); + if (!allow_multiple_reduction) { + throw std::runtime_error( + "Mooncake ElasticBuffer currently supports only " + "allow_multiple_reduction=true"); + } + if (num_buffer_bytes == 0) { + num_buffer_bytes = calculate_buffer_size( + num_ranks, num_max_tokens_per_rank, hidden, num_topk, + use_fp8_dispatch, allow_hybrid_mode, allow_multiple_reduction); + } + native_buffer_ = + std::make_unique(rank, num_ranks, num_buffer_bytes); + host_workspace_bytes_ = elastic_workspace_num_bytes(); + CUDA_CHECK(cudaHostAlloc(&host_workspace_, host_workspace_bytes_, + cudaHostAllocMapped)); + CUDA_CHECK( + cudaHostGetDevicePointer(&mapped_host_workspace_, host_workspace_, 0)); + std::memset(host_workspace_, 0, host_workspace_bytes_); +} + +MooncakeElasticBuffer::~MooncakeElasticBuffer() { + if (host_workspace_ != nullptr) { + cudaFreeHost(host_workspace_); + host_workspace_ = nullptr; + mapped_host_workspace_ = nullptr; + } +} + +int64_t MooncakeElasticBuffer::calculate_buffer_size( + int num_ranks, int64_t num_max_tokens_per_rank, int64_t hidden, + int64_t num_topk, bool use_fp8_dispatch, bool allow_hybrid_mode, + bool allow_multiple_reduction) { + num_topk = std::max(1, num_topk); + const int64_t dtype_bytes = use_fp8_dispatch ? 1 : 2; + const int64_t scale_bytes = + use_fp8_dispatch ? ceil_div_i64(hidden, 128) * 4 : 0; + const int64_t token_bytes = + align_i64(hidden * dtype_bytes, 32) + align_i64(scale_bytes, 32); + const int64_t metadata_bytes = align_i64( + num_topk * (sizeof(int) + sizeof(float)) + (1 + num_topk) * sizeof(int), + 32); + const int64_t per_slot_bytes = token_bytes + metadata_bytes; + const int64_t dispatch_bytes = + num_ranks * num_max_tokens_per_rank * num_topk * per_slot_bytes * 2; + const int64_t combine_factor = allow_multiple_reduction ? 3 : 4; + const int64_t combine_bytes = dispatch_bytes * combine_factor; + const int64_t hybrid_factor = allow_hybrid_mode && num_ranks > 1 ? 2 : 1; + return elastic_workspace_num_bytes() + elastic_atomic_scratch_num_bytes() + + hybrid_factor * (dispatch_bytes + combine_bytes); +} + +std::tuple MooncakeElasticBuffer::get_physical_domain_size() const { + return {topology_.num_rdma_ranks, topology_.num_nvlink_ranks}; +} + +std::tuple MooncakeElasticBuffer::get_logical_domain_size() const { + return {topology_.num_scaleout_ranks, topology_.num_scaleup_ranks}; +} + +int MooncakeElasticBuffer::get_theoretical_num_sms(int num_experts, + int num_topk) const { + int device = 0; + cudaGetDevice(&device); + cudaDeviceProp prop{}; + cudaGetDeviceProperties(&prop, device); + if (config_.prefer_overlap_with_compute) { + return std::max(1, std::min(24, prop.multiProcessorCount / 4)); + } + return std::max(1, std::min({40, prop.multiProcessorCount / 2, + std::max(1, num_experts * num_topk)})); +} + +ElasticDispatchOutput MooncakeElasticBuffer::dispatch( + const torch::Tensor& x, const std::optional& sf, + const torch::Tensor& topk_idx, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_experts, int num_max_tokens_per_rank, + int expert_alignment, int num_sms, bool do_expand, bool do_cpu_sync, + bool async_with_compute_stream, + const std::optional& cached_handle) { + EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous()); + const bool use_sf = sf.has_value(); + if (use_sf) { + EP_HOST_ASSERT(x.element_size() == 1); + EP_HOST_ASSERT(sf->dim() == 2 && sf->is_cuda()); + EP_HOST_ASSERT(sf->scalar_type() == torch::kFloat32 || + sf->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(sf->size(0) == x.size(0)); + } else { + EP_HOST_ASSERT(!config_.use_fp8_dispatch); + EP_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + } + EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous()); + EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64); + EP_HOST_ASSERT(x.size(0) == topk_idx.size(0)); + EP_HOST_ASSERT(num_experts % topology_.num_ranks == 0); + + const int num_tokens = static_cast(x.size(0)); + const int hidden = static_cast(x.size(1)); + const int num_topk = static_cast(topk_idx.size(1)); + const int num_sf_packs = use_sf ? static_cast(sf->size(1)) : 0; + const int sf_token_stride = use_sf ? static_cast(sf->stride(0)) : 0; + const int sf_hidden_stride = use_sf ? static_cast(sf->stride(1)) : 0; + const int num_local_experts = num_experts / topology_.num_ranks; + // The copy epilogue uses `kNumMaxTokensPerRank * kNumRanks` as the + // no-CPU-sync sentinel and then reads the real local receive count from the + // GPU prefix-sum tensor. In hybrid mode each scale-up peer may receive + // tokens forwarded from every scale-out rank, so the conservative output + // capacity and sentinel must cover the full logical world, not just the + // intra-node scale-up domain. + const int num_recv_tokens = num_max_tokens_per_rank * topology_.num_ranks; + const int num_smem_bytes = device_smem_bytes(); + const int num_channels_per_sm = 1; + const int num_channels = num_sms * num_channels_per_sm; + const bool cached_mode = cached_handle.has_value(); + const bool use_hybrid = topology_.num_scaleout_ranks != 1; + const int hybrid_channels = use_hybrid ? hybrid_num_channels(num_sms) : 0; + const int hybrid_max_tokens_per_channel = + use_hybrid ? hybrid_num_max_tokens_per_channel(num_max_tokens_per_rank, + num_sms) + : 0; + if (cached_mode) { + const auto& handle = cached_handle.value(); + EP_HOST_ASSERT(!handle.do_expand && !do_expand); + EP_HOST_ASSERT(handle.num_experts == num_experts); + EP_HOST_ASSERT(handle.expert_alignment == expert_alignment); + EP_HOST_ASSERT(handle.num_max_tokens_per_rank == + num_max_tokens_per_rank); + EP_HOST_ASSERT(handle.num_sms == num_sms); + if (use_hybrid) { + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.dim() == 4); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(0) == + hybrid_channels); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(1) == + topology_.num_scaleout_ranks); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(2) == + hybrid_max_tokens_per_channel); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(3) == num_topk); + EP_HOST_ASSERT(handle.token_metadata_at_forward.has_value()); + EP_HOST_ASSERT(handle.channel_linked_list.has_value()); + } else { + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.dim() == 2); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(0) == num_tokens); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(1) == num_topk); + } + } + + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = native_buffer_->comm_stream; + stream_wait(launch_stream, compute_stream); + + const int64_t timeout_cycles = + config_.num_gpu_timeout_secs < 0 + ? -1 + : static_cast(native_buffer_->clock_rate_khz) * + static_cast(config_.num_gpu_timeout_secs) * 1000; + auto launch_ctx = make_launch_context( + *native_buffer_, topology_, mapped_host_workspace_, timeout_cycles); + + auto psum_num_recv_tokens_per_scaleup_rank = + cached_mode ? cached_handle->psum_num_recv_tokens_per_scaleup_rank + : torch::empty({topology_.num_scaleup_ranks}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())); + auto psum_num_recv_tokens_per_expert = + cached_mode + ? cached_handle->psum_num_recv_tokens_per_expert + : torch::empty({num_local_experts + 1}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())); + auto dst_buffer_slot_idx = + cached_mode + ? cached_handle->dst_buffer_slot_idx + : (use_hybrid ? torch::empty( + {hybrid_channels, topology_.num_scaleout_ranks, + hybrid_max_tokens_per_channel, num_topk}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())) + : torch::empty({num_tokens, num_topk}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device()))); + std::optional token_metadata_at_forward = std::nullopt; + std::optional channel_linked_list = std::nullopt; + if (use_hybrid) { + if (cached_mode) { + token_metadata_at_forward = + cached_handle->token_metadata_at_forward; + channel_linked_list = cached_handle->channel_linked_list; + } else { + const int forward_metadata_dims = 2 + num_topk * 2; + token_metadata_at_forward = torch::empty( + {hybrid_channels, + topology_.num_scaleout_ranks * hybrid_max_tokens_per_channel + + 1, + forward_metadata_dims}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + channel_linked_list = torch::empty( + {hybrid_channels, + topology_.num_scaleout_ranks * hybrid_max_tokens_per_channel + + 1, + topology_.num_scaleup_ranks}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + } + } + std::optional deterministic_rank_count_buffer = std::nullopt; +#ifdef MOONCAKE_EP_USE_MUSA + // MUSA non-hybrid dispatch always runs + // launch_musa_elastic_prepare_dispatch(), which assigns slots and publishes + // counts without cooperative grid sync. + const bool run_deterministic_prologue = false; +#else + const bool run_deterministic_prologue = + config_.deterministic && !cached_mode && !use_hybrid; +#endif + if (run_deterministic_prologue) { + deterministic_rank_count_buffer = torch::empty( + {num_sms, topology_.num_scaleup_ranks}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + launch_elastic_dispatch_deterministic_prologue( + topk_idx.data_ptr(), + deterministic_rank_count_buffer.value().data_ptr(), + dst_buffer_slot_idx.data_ptr(), num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, + topology_.scaleup_rank_idx, topology_.num_scaleup_ranks, num_sms, + num_smem_bytes, launch_stream.stream()); + } + + launch_mooncake_elastic_dispatch( + x.data_ptr(), use_sf ? const_cast(sf->data_ptr()) : nullptr, + const_cast(topk_idx.data_ptr()), + topk_weights.has_value() + ? const_cast(topk_weights->data_ptr()) + : nullptr, + nullptr, nullptr, psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + psum_num_recv_tokens_per_expert.data_ptr(), + dst_buffer_slot_idx.data_ptr(), + token_metadata_at_forward.has_value() + ? token_metadata_at_forward->data_ptr() + : nullptr, + num_tokens, num_max_tokens_per_rank, hidden, + static_cast(x.element_size()), num_sf_packs, sf_token_stride, + sf_hidden_stride, num_experts, num_topk, expert_alignment, num_sms, + use_hybrid ? kElasticHybridChannelsPerSm : num_channels_per_sm, + num_smem_bytes, cached_mode, config_.deterministic, false, launch_ctx, + launch_stream.stream()); + + const int num_recv_output_capacity = + do_expand ? num_recv_tokens * num_topk : num_recv_tokens; + auto recv_x = torch::empty({num_recv_output_capacity, hidden}, x.options()); + auto recv_x_scales = std::optional(); + void* recv_x_scales_ptr = nullptr; + int recv_sf_token_stride = 0; + int recv_sf_hidden_stride = 0; + if (use_sf) { + recv_x_scales = torch::empty({num_recv_output_capacity, num_sf_packs}, + sf->options()); + recv_x_scales_ptr = recv_x_scales->data_ptr(); + recv_sf_token_stride = static_cast(recv_x_scales->stride(0)); + recv_sf_hidden_stride = static_cast(recv_x_scales->stride(1)); + } + auto recv_topk_idx = + torch::empty({num_recv_tokens, num_topk}, topk_idx.options()); + auto recv_topk_weights = std::optional(); + float* recv_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + recv_topk_weights = do_expand + ? torch::empty({num_recv_output_capacity}, + topk_weights->options()) + : torch::empty({num_recv_tokens, num_topk}, + topk_weights->options()); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + auto recv_src_metadata = torch::empty( + {num_recv_tokens, num_topk + 2}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + auto handle_psum_num_recv_tokens_per_expert = + do_expand + ? psum_num_recv_tokens_per_expert.slice(0, 0, num_local_experts) + : psum_num_recv_tokens_per_expert.slice(0, 1, + num_local_experts + 1); + auto epilogue_psum_num_recv_tokens_per_expert = + do_expand ? psum_num_recv_tokens_per_expert + : handle_psum_num_recv_tokens_per_expert; + + launch_mooncake_elastic_dispatch_copy_epilogue( + recv_x.data_ptr(), recv_x_scales_ptr, recv_topk_idx.data_ptr(), + recv_topk_weights_ptr, recv_src_metadata.data_ptr(), + channel_linked_list.has_value() ? channel_linked_list->data_ptr() + : nullptr, + num_recv_tokens, num_max_tokens_per_rank, hidden, + static_cast(x.element_size()), num_sf_packs, recv_sf_token_stride, + recv_sf_hidden_stride, num_experts, num_topk, num_sms, num_smem_bytes, + use_hybrid ? hybrid_channels : num_channels, do_expand, cached_mode, + launch_ctx, psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + epilogue_psum_num_recv_tokens_per_expert.data_ptr(), + launch_stream.stream()); + + if (do_cpu_sync || !async_with_compute_stream) { + stream_wait(compute_stream, launch_stream); + } + std::optional event = std::nullopt; + if (async_with_compute_stream) { + event = EventHandle(launch_stream); + } + + std::vector num_recv_tokens_per_expert_list; + int actual_num_recv_tokens = num_recv_tokens; + int actual_num_output_tokens = num_recv_tokens; + if (do_cpu_sync) { + auto scaleup_psum_cpu = psum_num_recv_tokens_per_scaleup_rank.cpu(); + auto expert_psum_cpu = psum_num_recv_tokens_per_expert.cpu(); + const auto* scaleup_psum = scaleup_psum_cpu.data_ptr(); + const auto* expert_psum = expert_psum_cpu.data_ptr(); + actual_num_recv_tokens = scaleup_psum[topology_.num_scaleup_ranks - 1]; + EP_HOST_ASSERT(actual_num_recv_tokens >= 0 && + actual_num_recv_tokens <= num_recv_tokens); + actual_num_output_tokens = actual_num_recv_tokens; + + num_recv_tokens_per_expert_list.reserve(num_local_experts); + const auto align_count = [expert_alignment](int value) { + return ((value + expert_alignment - 1) / expert_alignment) * + expert_alignment; + }; + if (do_expand) { + int previous_psum = 0; + for (int i = 0; i < num_local_experts; ++i) { + const int count = expert_psum[i] - align_count(previous_psum); + EP_HOST_ASSERT(count >= 0); + num_recv_tokens_per_expert_list.push_back(count); + previous_psum = expert_psum[i]; + } + actual_num_output_tokens = + num_local_experts == 0 ? 0 : expert_psum[num_local_experts - 1]; + } else { + for (int i = 0; i < num_local_experts; ++i) { + const int count = expert_psum[i + 1] - expert_psum[i]; + EP_HOST_ASSERT(count >= 0); + num_recv_tokens_per_expert_list.push_back(count); + } + } + EP_HOST_ASSERT(actual_num_output_tokens >= 0 && + actual_num_output_tokens <= recv_x.size(0)); + + recv_x = recv_x.slice(0, 0, actual_num_output_tokens); + if (recv_x_scales.has_value()) { + recv_x_scales = + recv_x_scales->slice(0, 0, actual_num_output_tokens); + } + recv_topk_idx = recv_topk_idx.slice(0, 0, actual_num_recv_tokens); + if (recv_topk_weights.has_value()) { + recv_topk_weights = + recv_topk_weights->slice(0, 0, actual_num_output_tokens); + } + recv_src_metadata = + recv_src_metadata.slice(0, 0, actual_num_recv_tokens); + } + + ElasticNativeHandle handle; + handle.do_expand = do_expand; + handle.num_experts = num_experts; + handle.expert_alignment = expert_alignment; + handle.num_max_tokens_per_rank = num_max_tokens_per_rank; + handle.num_sms = num_sms; + handle.topk_idx = cached_mode ? cached_handle->topk_idx : topk_idx.clone(); + handle.psum_num_recv_tokens_per_expert = + handle_psum_num_recv_tokens_per_expert; + handle.psum_num_recv_tokens_per_scaleup_rank = + psum_num_recv_tokens_per_scaleup_rank; + handle.recv_src_metadata = recv_src_metadata; + handle.recv_layout_range = torch::empty( + {0}, torch::TensorOptions().dtype(torch::kInt64).device(x.device())); + handle.dst_buffer_slot_idx = dst_buffer_slot_idx; + handle.token_metadata_at_forward = token_metadata_at_forward; + handle.channel_linked_list = channel_linked_list; + handle.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list; + + ElasticDispatchOutput output; + output.recv_x = recv_x; + output.recv_x_scales = recv_x_scales; + output.recv_topk_idx = recv_topk_idx; + output.recv_topk_weights = recv_topk_weights; + output.handle = handle; + output.event = event; + return output; +} + +ElasticCombineOutput MooncakeElasticBuffer::combine( + const torch::Tensor& x, const ElasticNativeHandle& handle, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_sms, bool async_with_compute_stream, + const std::optional& out) { + EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous()); + EP_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + torch::Tensor weights = topk_weights.value_or(torch::Tensor()); + if (!weights.defined()) { + weights = torch::ones( + handle.topk_idx.sizes(), + torch::TensorOptions().dtype(torch::kFloat32).device(x.device())); + } + const int hidden = static_cast(x.size(1)); + const int num_topk = static_cast(handle.topk_idx.size(1)); + const int num_combined_tokens = static_cast(handle.topk_idx.size(0)); + const int num_smem_bytes = device_smem_bytes(); + const int num_channels = std::max(1, num_sms); + const bool use_hybrid = topology_.num_scaleout_ranks != 1; + const int hybrid_channels = use_hybrid ? hybrid_num_channels(num_sms) : 0; + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = native_buffer_->comm_stream; + stream_wait(launch_stream, compute_stream); + const int64_t timeout_cycles = + config_.num_gpu_timeout_secs < 0 + ? -1 + : static_cast(native_buffer_->clock_rate_khz) * + static_cast(config_.num_gpu_timeout_secs) * 1000; + auto launch_ctx = make_launch_context( + *native_buffer_, topology_, mapped_host_workspace_, timeout_cycles); + auto psum_num_recv_tokens_per_scaleup_rank = + handle.psum_num_recv_tokens_per_scaleup_rank; + void* reduce_buffer = launch_mooncake_elastic_combine( + x.data_ptr(), weights.data_ptr(), + const_cast(handle.recv_src_metadata.data_ptr()), + psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + handle.token_metadata_at_forward.has_value() + ? handle.token_metadata_at_forward->data_ptr() + : nullptr, + handle.channel_linked_list.has_value() + ? handle.channel_linked_list->data_ptr() + : nullptr, + static_cast(x.size(0)), handle.num_max_tokens_per_rank, hidden, + handle.num_experts, num_topk, num_sms, num_smem_bytes, + use_hybrid ? hybrid_channels : num_channels, handle.do_expand, + config_.allow_multiple_reduction, launch_ctx, launch_stream.stream()); + + torch::Tensor combined_x = + out.has_value() + ? out.value() + : torch::empty({num_combined_tokens, hidden}, x.options()); + launch_mooncake_elastic_combine_reduce_epilogue( + combined_x.data_ptr(), weights.data_ptr(), + const_cast(handle.topk_idx.data_ptr()), + num_combined_tokens, handle.num_max_tokens_per_rank, hidden, + handle.num_experts, num_topk, reduce_buffer, nullptr, nullptr, num_sms, + num_smem_bytes, handle.do_expand, config_.allow_multiple_reduction, + launch_ctx, launch_stream.stream()); + + if (!async_with_compute_stream) { + stream_wait(compute_stream, launch_stream); + } + std::optional event = std::nullopt; + if (async_with_compute_stream) event = EventHandle(launch_stream); + (void)active_ranks; + + ElasticCombineOutput output; + output.combined_x = combined_x; + output.combined_topk_weights = std::nullopt; + output.event = event; + return output; +} + +ElasticTopology MooncakeElasticBuffer::discover_topology( + int rank, int num_ranks, bool allow_hybrid_mode) { + int device_count = 1; + cudaGetDeviceCount(&device_count); + int num_local_ranks = + getenv_int("MOONCAKE_EP_NUM_LOCAL_RANKS", + std::max(1, std::min(num_ranks, device_count))); + num_local_ranks = std::max(1, std::min(num_local_ranks, num_ranks)); + + ElasticTopology topology; + topology.rank_idx = rank; + topology.num_ranks = num_ranks; + topology.num_rdma_ranks = + static_cast(ceil_div_i64(num_ranks, num_local_ranks)); + topology.num_nvlink_ranks = num_local_ranks; + if (allow_hybrid_mode && topology.num_rdma_ranks > 1) { + topology.num_scaleout_ranks = topology.num_rdma_ranks; + topology.num_scaleup_ranks = topology.num_nvlink_ranks; + topology.hybrid_enabled = true; + } else { + topology.num_scaleout_ranks = 1; + topology.num_scaleup_ranks = num_ranks; + topology.hybrid_enabled = false; + } + topology.scaleout_rank_idx = rank / topology.num_scaleup_ranks; + topology.scaleup_rank_idx = rank % topology.num_scaleup_ranks; + return topology; +} + +} // namespace mooncake diff --git a/mooncake-ep/src/mooncake_ep_elastic_kernel.cu b/mooncake-ep/src/mooncake_ep_elastic_kernel.cu new file mode 100644 index 0000000000..a5d3e383ff --- /dev/null +++ b/mooncake-ep/src/mooncake_ep_elastic_kernel.cu @@ -0,0 +1,973 @@ +// clang-format off + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +constexpr int kElasticNumNotifyWarps = 4; +#ifdef MOONCAKE_EP_USE_MUSA +constexpr int kElasticNumDispatchWarps = 4; +constexpr int kElasticNumEpilogueWarps = 4; +#else +constexpr int kElasticNumDispatchWarps = 8; +constexpr int kElasticNumEpilogueWarps = 8; +#endif +constexpr int kElasticNumHybridScaleoutWarps = 4; +constexpr int kElasticNumHybridForwardWarps = 4; +constexpr int kElasticNumHybridScaleupWarps = 4; +constexpr int kElasticNumQPs = MAX_QP_COUNT; +constexpr int64_t kElasticTimeoutCycles = NUM_TIMEOUT_CYCLES; + +inline int ceil_div(int x, int y) { return (x + y - 1) / y; } + +inline int hybrid_num_channels(int num_sms) { + return num_sms * kElasticNumHybridForwardWarps; +} + +inline void* hybrid_combine_reduce_buffer_ptr(void* buffer, int hidden, + int num_topk, + int num_max_tokens_per_rank, + int num_scaleout_ranks, + int num_scaleup_ranks, + bool allow_multiple_reduction) { + const int num_tokens_in_scaleup_layout = + allow_multiple_reduction && num_scaleup_ranks <= num_topk + ? num_scaleup_ranks + : num_topk; + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, num_topk, false); + const auto scaleup_buffer = elastic::layout::BufferLayout( + token_layout, num_tokens_in_scaleup_layout, + num_scaleout_ranks * num_max_tokens_per_rank, buffer); + return scaleup_buffer.get_buffer_end_ptr(); +} + +inline int dispatch_smem_bytes(int hidden, int elem_size, int num_sf_packs, + int num_topk, int num_ranks, + int num_experts, int num_notify_warps, + int num_dispatch_warps) { + const int notify_smem_bytes = num_notify_warps == 0 + ? 0 + : elastic::math::align(num_ranks + num_experts, num_notify_warps * 32) * + static_cast(sizeof(int)); + const auto token_layout = + elastic::layout::TokenLayout(hidden * elem_size, + num_sf_packs * sizeof(sf_pack_t), + num_topk, true); + return notify_smem_bytes + + num_dispatch_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int dispatch_epilogue_smem_bytes(int hidden, int elem_size, + int num_sf_packs, int num_topk, + int num_warps) { + const auto token_layout = + elastic::layout::TokenLayout(hidden * elem_size, + num_sf_packs * sizeof(sf_pack_t), + num_topk, true); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int combine_smem_bytes(int hidden, int num_topk, int num_warps) { + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, num_topk, false); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int combine_epilogue_smem_bytes(int hidden, int num_warps) { + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, 0, false); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline device::CommCtx make_comm_ctx(const ElasticLaunchContext& ctx) { + device::CommCtx comm_ctx{}; + comm_ctx.rank = ctx.rank; + comm_ctx.p2p.available = ctx.nvlink_available; + comm_ctx.p2p.peer_ptrs = ctx.ipc_peer_ptrs; + comm_ctx.p2p.local_base = ctx.gdr_buffer; + comm_ctx.ibgda.qp_devctxs = + reinterpret_cast(ctx.qp_devctxs); + comm_ctx.ibgda.raddrs = reinterpret_cast(ctx.raddrs); + comm_ctx.ibgda.rkeys = reinterpret_cast(ctx.rkeys); + comm_ctx.ibgda.local_atomic_base = ctx.rdma_send_signal_buffer; + comm_ctx.ibgda.remote_atomic_base = ctx.rdma_recv_signal_buffer; + return comm_ctx; +} + +#ifdef MOONCAKE_EP_USE_MUSA + +// MUSA currently cannot rely on CUDA-style cooperative grid synchronization for +// the official elastic dispatch notify/prologue. These prepare kernels keep the +// dispatch algorithm semantics unchanged (slot assignment, count publication and +// prefix-sum generation), but split that prologue into ordinary launches with an +// explicit scale-up barrier so peer count writes cannot race with local clears. +// Payload movement still goes through the common dispatch kernel and Mooncake +// Device API transport primitives; this is only a no-cooperative-grid-sync +// metadata preparation fallback. + +__global__ void musa_elastic_prepare_init_kernel(void* workspace, + int num_scaleup_ranks, + int num_experts) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = blockDim.x * gridDim.x; + for (int i = tid; i < num_scaleup_ranks + num_experts; i += stride) { + layout.get_scaleup_rank_expert_count_ptr()[i] = 0; + layout.get_scaleup_rank_expert_count_ptr()[i] = 0; + } + for (int i = tid; i < num_scaleup_ranks; i += stride) { + layout.get_scaleup_atomic_sender_counter()[i] = 0; + } +} + +__global__ void musa_elastic_prepare_clear_barrier_kernel( + device::CommCtx comm_ctx, void* workspace, int rank_idx, + int num_scaleup_ranks, int num_experts, int64_t timeout_cycles) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const auto gin = elastic::transport::MooncakeGin( + comm_ctx, 0, 0, 1, 0, rank_idx, num_scaleup_ranks, num_scaleup_ranks); + constexpr int kTag = elastic::comm::kDeviceBarrierTag; + const int status = + static_cast((*layout.get_nvl_barrier_counter_ptr(kTag)) & 3); + const int phase = status & 1; + const int sign = status >> 1; + const int* base_signal = layout.get_nvl_barrier_signal_ptr(kTag, phase); + + if (threadIdx.x < num_scaleup_ranks) { + auto* dst_ptr = const_cast(base_signal) + rank_idx; + gin.red_add_rel( + dst_ptr, sign ? -1 : 1, threadIdx.x); + } + __syncthreads(); + + if (threadIdx.x == 0) { + atomicAdd(layout.get_nvl_barrier_counter_ptr(kTag), 1ULL); + const int target = sign ? 0 : num_scaleup_ranks; + const auto start_clock = clock64(); + while (true) { + int sum = 0; + for (int i = 0; i < num_scaleup_ranks; ++i) { + sum += elastic::ptx::ld_acquire_sys( + const_cast(base_signal) + i); + } + if (sum == target) break; + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf( + "MUSA prepare clear barrier timeout, rank=%d sum=%d " + "target=%d\n", + rank_idx, sum, target); + break; + } + } + } +} + +__global__ void musa_elastic_assign_slots_kernel( + const int64_t* topk_idx, int* dst_buffer_slot_idx, void* workspace, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int num_scaleup_ranks, int rank_idx) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + const int token_stride = blockDim.x * gridDim.x; + for (int token_idx = blockIdx.x * blockDim.x + threadIdx.x; + token_idx < num_tokens; token_idx += token_stride) { + int seen_ranks[32]; + int num_seen = 0; + for (int k = 0; k < num_topk; ++k) { + dst_buffer_slot_idx[token_idx * num_topk + k] = -1; + } + for (int k = 0; k < num_topk; ++k) { + const int expert_idx = + static_cast(topk_idx[token_idx * num_topk + k]); + if (expert_idx < 0) continue; + const int dst_rank = expert_idx / num_experts_per_rank; + bool duplicate_rank = false; + for (int i = 0; i < num_seen; ++i) + duplicate_rank |= (seen_ranks[i] == dst_rank); + if (!duplicate_rank) { + seen_ranks[num_seen++] = dst_rank; + const int slot = atomicAdd( + layout.get_scaleup_atomic_sender_counter() + dst_rank, 1); + dst_buffer_slot_idx[token_idx * num_topk + k] = + rank_idx * num_max_tokens_per_rank + slot; + atomicAdd(reinterpret_cast( + layout.get_scaleup_rank_count_ptr() + + dst_rank), + 1ULL); + } + atomicAdd(reinterpret_cast( + layout.get_scaleup_expert_count_ptr() + + expert_idx), + 1ULL); + } + } +} + +__global__ void musa_elastic_publish_counts_kernel( + device::CommCtx comm_ctx, void* workspace, int rank_idx, + int num_scaleup_ranks, int num_experts) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + const auto gin = elastic::transport::MooncakeGin( + comm_ctx, 0, 0, 1, 0, rank_idx, num_scaleup_ranks, num_scaleup_ranks); + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = blockDim.x * gridDim.x; + for (int dst_rank = tid; dst_rank < num_scaleup_ranks; dst_rank += stride) { + const auto count = static_cast( + layout.get_scaleup_rank_count_ptr()[dst_rank]); + auto* dst = reinterpret_cast( + layout.get_scaleup_rank_count_ptr() + rank_idx); + const auto encoded_count = + elastic::math::encode_decode_positive(count); + if (dst_rank == rank_idx) { + elastic::ptx::st_release_sys(dst, encoded_count); + } else { + gin.put_value( + dst, encoded_count, dst_rank, 0); + } + } + for (int expert_idx = tid; expert_idx < num_experts; expert_idx += stride) { + const int dst_rank = expert_idx / num_experts_per_rank; + const int local_expert_idx = expert_idx % num_experts_per_rank; + const auto count = static_cast( + layout.get_scaleup_expert_count_ptr()[expert_idx]); + auto* dst = reinterpret_cast( + layout.get_scaleup_expert_count_ptr() + + rank_idx * num_experts_per_rank + local_expert_idx); + const auto encoded_count = + elastic::math::encode_decode_positive(count); + if (dst_rank == rank_idx) { + elastic::ptx::st_release_sys(dst, encoded_count); + } else { + gin.put_value( + dst, encoded_count, dst_rank, 0); + } + } +} + +__global__ void musa_elastic_wait_prefix_counts_kernel( + void* workspace, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int rank_idx, int num_scaleup_ranks, + int num_experts, int expert_alignment, int64_t timeout_cycles) { + (void)rank_idx; + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + if (threadIdx.x == 0 && blockIdx.x == 0) { + int psum = 0; + for (int src_rank = 0; src_rank < num_scaleup_ranks; ++src_rank) { + auto* ptr = layout.get_scaleup_rank_count_ptr() + src_rank; + const auto start_clock = clock64(); + auto* word_ptr = reinterpret_cast(ptr); + int count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + while (!elastic::math::is_decoded_positive_ready(count)) { + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf("MUSA prepare rank-count timeout, self=%d src=%d decoded=%d\n", + rank_idx, src_rank, count); + count = 0; + break; + } + count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + } + *ptr = 0; + psum += count; + psum_num_recv_tokens_per_scaleup_rank[src_rank] = psum; + } + psum = 0; + psum_num_recv_tokens_per_expert[0] = 0; + for (int expert_idx = 0; expert_idx < num_experts_per_rank; + ++expert_idx) { + int count = 0; + for (int src_rank = 0; src_rank < num_scaleup_ranks; ++src_rank) { + auto* ptr = layout.get_scaleup_expert_count_ptr() + + src_rank * num_experts_per_rank + expert_idx; + const auto start_clock = clock64(); + auto* word_ptr = reinterpret_cast(ptr); + int encoded_count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + while (!elastic::math::is_decoded_positive_ready(encoded_count)) { + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf("MUSA prepare expert-count timeout, self=%d src=%d expert=%d decoded=%d\n", + rank_idx, src_rank, expert_idx, encoded_count); + encoded_count = 0; + break; + } + encoded_count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + } + count += encoded_count; + *ptr = 0; + } + psum += elastic::math::align(count, expert_alignment); + psum_num_recv_tokens_per_expert[expert_idx + 1] = psum; + } + } +} + +void launch_musa_elastic_prepare_dispatch( + const int64_t* topk_idx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int expert_alignment, const device::CommCtx& comm_ctx, + const ElasticLaunchContext& ctx, cudaStream_t stream) { + constexpr int kThreads = 256; + const int blocks = std::max(1, std::min(128, ceil_div(num_tokens, kThreads))); + musa_elastic_prepare_init_kernel<<>>( + ctx.workspace, ctx.num_scaleup_ranks, num_experts); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + // Each rank clears its local receive-count slots in the init kernel above. + // Without a cross-rank phase boundary, a fast peer may publish counts into + // this rank while its init kernel is still clearing the same slots, losing + // the peer write. CUDA's cooperative prologue gets this ordering from grid + // synchronization; MUSA needs an explicit scale-up barrier before publish. + musa_elastic_prepare_clear_barrier_kernel<<<1, kThreads, 0, stream>>>( + comm_ctx, ctx.workspace, ctx.scaleup_rank_idx, ctx.num_scaleup_ranks, + num_experts, ctx.timeout_cycles); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_assign_slots_kernel<<>>( + topk_idx, dst_buffer_slot_idx, ctx.workspace, num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, ctx.num_scaleup_ranks, + ctx.scaleup_rank_idx); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_publish_counts_kernel<<>>( + comm_ctx, ctx.workspace, ctx.scaleup_rank_idx, ctx.num_scaleup_ranks, + num_experts); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_wait_prefix_counts_kernel<<<1, 1, 0, stream>>>( + ctx.workspace, psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, ctx.scaleup_rank_idx, + ctx.num_scaleup_ranks, num_experts, expert_alignment, + ctx.timeout_cycles); + CUDA_RUNTIME_CHECK(cudaGetLastError()); +} + +#endif + +template +void launch_cooperative(Kernel kernel, int num_sms, int num_threads, + int smem_bytes, cudaStream_t stream, Args... args) { +#ifndef MOONCAKE_EP_USE_MUSA + CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)); +#endif +#ifdef MOONCAKE_EP_USE_MUSA + kernel<<>>(args...); + CUDA_RUNTIME_CHECK(cudaGetLastError()); +#else + cudaLaunchConfig_t cfg = {{num_sms, 1, 1}, {num_threads, 1, 1}, + static_cast(smem_bytes), stream, + nullptr, 0}; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeCooperative; + attr[0].val.cooperative = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&cfg, kernel, args...)); +#endif +} + +[[noreturn]] void unsupported_elastic_config(const char* op, int hidden, + int num_experts, int num_topk, + int num_max_tokens_per_rank, + int num_sms, int ranks) { + throw std::runtime_error( + std::string("Unsupported Mooncake elastic ") + op + + " static-template config: hidden=" + std::to_string(hidden) + + ", experts=" + std::to_string(num_experts) + + ", topk=" + std::to_string(num_topk) + + ", max_tokens=" + std::to_string(num_max_tokens_per_rank) + + ", num_sms=" + std::to_string(num_sms) + + ", ranks=" + std::to_string(ranks)); +} + +} // namespace + +void launch_elastic_dispatch_deterministic_prologue( + const int64_t* topk_idx, int* rank_count_buffer, int* dst_buffer_slot_idx, + int num_tokens, + int num_max_tokens_per_rank, int num_experts, int num_topk, + int scaleup_rank_idx, int num_scaleup_ranks, int num_sms, + int num_smem_bytes, cudaStream_t stream) { + constexpr int kNumWarps = kElasticNumEpilogueWarps; + constexpr int kNumThreads = kNumWarps * 32; + const int smem_bytes = (1 + 2 * kNumWarps) * num_scaleup_ranks * sizeof(int); + (void)num_smem_bytes; + +#define LAUNCH_PROLOGUE(HIDDEN, EXPERTS, TOPK, MAXTOK, SMS, RANKS) \ + do { \ + auto kernel = elastic::dispatch_deterministic_prologue_impl< \ + SMS, kNumWarps, RANKS, MAXTOK, EXPERTS, TOPK>; \ + launch_cooperative(kernel, SMS, kNumThreads, smem_bytes, stream, \ + const_cast(topk_idx), rank_count_buffer, \ + dst_buffer_slot_idx, num_tokens, \ + scaleup_rank_idx); \ + } while (false) + +#define TRY_PROLOGUE(H, E, K, M, S, R) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + num_scaleup_ranks == R) { \ + LAUNCH_PROLOGUE(H, E, K, M, S, R); \ + return; \ + } + + const int hidden = 0; + (void)hidden; +#ifdef MOONCAKE_EP_USE_MUSA + // Keep the MUSA compile set intentionally small while validating the + // native elastic scale-up path; MUSA non-hybrid dispatch prepares slots in + // a separate kernel and does not call this CUDA cooperative prologue. + TRY_PROLOGUE(0, 256, 8, 128, 24, 2); + TRY_PROLOGUE(0, 256, 8, 128, 24, 8); +#else + // Common production MoE shapes; hidden is irrelevant for this prologue. + TRY_PROLOGUE(0, 256, 8, 128, 24, 8); + TRY_PROLOGUE(0, 256, 8, 128, 24, 2); +#endif + +#undef TRY_PROLOGUE +#undef LAUNCH_PROLOGUE + unsupported_elastic_config("deterministic_prologue", 0, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + num_scaleup_ranks); +} + +void launch_mooncake_elastic_dispatch( + void* x, void* sf, int64_t* topk_idx, float* topk_weights, + int64_t* copied_topk_idx, int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int* token_metadata_at_forward, int num_tokens, + int num_max_tokens_per_rank, int hidden, int elem_size, int num_sf_packs, + int sf_token_stride, int sf_hidden_stride, int num_experts, int num_topk, + int expert_alignment, int num_sms, int num_channels_per_sm, + int num_smem_bytes, bool cached_mode, bool deterministic, + bool do_cpu_sync, const ElasticLaunchContext& ctx, cudaStream_t stream) { +#ifdef MOONCAKE_EP_USE_MUSA + const bool musa_use_prepared_slots = !cached_mode && ctx.num_scaleout_ranks == 1; +#else + const bool musa_use_prepared_slots = false; +#endif + const bool effective_cached_mode = cached_mode || musa_use_prepared_slots; + const int num_notify_warps = effective_cached_mode ? 0 : kElasticNumNotifyWarps; + const int num_dispatch_warps = kElasticNumDispatchWarps; + const int num_threads = (num_notify_warps + num_dispatch_warps) * 32; + const int smem_bytes = std::max( + num_smem_bytes, + dispatch_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + ctx.num_scaleup_ranks, num_experts, + num_notify_warps, num_dispatch_warps)); + const bool reuse_slot_indices = effective_cached_mode || deterministic; + const auto comm_ctx = make_comm_ctx(ctx); + (void)num_channels_per_sm; + +#ifdef MOONCAKE_EP_USE_MUSA + if (musa_use_prepared_slots) { + launch_musa_elastic_prepare_dispatch( + topk_idx, psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, dst_buffer_slot_idx, num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, expert_alignment, + comm_ctx, ctx, stream); + } +#endif + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { + const bool hybrid_reuse_slot_indices = cached_mode; + const int hybrid_dispatch_warps = + kElasticNumHybridScaleoutWarps + kElasticNumHybridForwardWarps; + const int hybrid_threads = + (num_notify_warps + hybrid_dispatch_warps) * 32; + const int hybrid_smem_bytes = std::max( + num_smem_bytes, + dispatch_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + ctx.num_scaleout_ranks * ctx.num_scaleup_ranks, + num_experts, num_notify_warps, + hybrid_dispatch_warps)); + +#define LAUNCH_HYBRID_DISPATCH(HB, SFP, E, K, M, S, SO, SU) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + if (cached_mode) { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, true, S, 0, kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } else if (hybrid_reuse_slot_indices) { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, true, S, kElasticNumNotifyWarps, \ + kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } else { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, false, S, kElasticNumNotifyWarps, \ + kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } \ + } while (false) + +#define TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + elem_size == EL && num_sf_packs == SFP && expert_alignment == 1 && \ + !do_cpu_sync) { \ + LAUNCH_HYBRID_DISPATCH((H) * (EL), SFP, E, K, M, S, SO, SU); \ + return; \ + } + +#define TRY_HYBRID_DISPATCH(H, E, K, M, S, SO, SU) \ + TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, \ + static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, 1, (H) / 128) + +#define TRY_HYBRID_DISPATCH_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_DISPATCH(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_DISPATCH(H, E, K, M, S, 2, 8) + + TRY_HYBRID_DISPATCH_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_DISPATCH_SHAPE +#undef TRY_HYBRID_DISPATCH +#undef TRY_HYBRID_DISPATCH_TYPED +#undef LAUNCH_HYBRID_DISPATCH + } +#endif + +#define LAUNCH_DISPATCH(HB, SFP, E, K, M, S, R) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + if (effective_cached_mode) { \ + auto kernel = elastic::dispatch_impl< \ + true, false, true, S, 0, kElasticNumDispatchWarps, R, \ + kHiddenBytes, kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } else if (reuse_slot_indices) { \ + auto kernel = elastic::dispatch_impl< \ + true, false, true, S, kElasticNumNotifyWarps, \ + kElasticNumDispatchWarps, R, kHiddenBytes, kNumSFPacks, M, E, K, 1, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } else { \ + auto kernel = elastic::dispatch_impl< \ + true, false, false, S, kElasticNumNotifyWarps, \ + kElasticNumDispatchWarps, R, kHiddenBytes, kNumSFPacks, M, E, K, 1, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } \ + } while (false) + +#define TRY_DISPATCH_TYPED(H, E, K, M, S, R, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && elem_size == EL && \ + num_sf_packs == SFP && expert_alignment == 1 && !do_cpu_sync) { \ + LAUNCH_DISPATCH((H) * (EL), SFP, E, K, M, S, R); \ + return; \ + } + +#define TRY_DISPATCH(H, E, K, M, S, R) \ + TRY_DISPATCH_TYPED(H, E, K, M, S, R, static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_DISPATCH_TYPED(H, E, K, M, S, R, 1, (H) / 128) + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_DISPATCH_TYPED(4096, 256, 8, 128, 24, 2, + static_cast(sizeof(nv_bfloat16)), 0); + TRY_DISPATCH_TYPED(4096, 256, 8, 128, 24, 8, + static_cast(sizeof(nv_bfloat16)), 0); +#else + TRY_DISPATCH(4096, 256, 8, 128, 24, 8); + TRY_DISPATCH(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_DISPATCH +#undef TRY_DISPATCH_TYPED +#undef LAUNCH_DISPATCH + unsupported_elastic_config("dispatch", hidden, num_experts, num_topk, + num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void launch_mooncake_elastic_dispatch_copy_epilogue( + void* recv_x, void* recv_sf, int64_t* recv_topk_idx, + float* recv_topk_weights, int* recv_src_metadata, + int* channel_linked_list, int num_recv_tokens, int num_max_tokens_per_rank, + int hidden, int elem_size, int num_sf_packs, int recv_sf_token_stride, + int recv_sf_hidden_stride, int num_experts, int num_topk, int num_sms, + int num_smem_bytes, int num_channels, bool do_expand, bool cached_mode, + const ElasticLaunchContext& ctx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, + dispatch_epilogue_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + kElasticNumEpilogueWarps)); + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { +#define LAUNCH_HYBRID_DISPATCH_EPILOGUE(HB, SFP, E, K, M, S, SO, SU, C) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + auto kernel = do_expand ? \ + elastic::dispatch_copy_epilogue_impl< \ + true, false, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + (cached_mode ? \ + elastic::dispatch_copy_epilogue_impl< \ + false, true, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + elastic::dispatch_copy_epilogue_impl< \ + false, false, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K>); \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + ctx.buffer, ctx.workspace, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, recv_x, \ + static_cast(recv_sf), \ + recv_topk_idx, recv_topk_weights, \ + recv_src_metadata, channel_linked_list, \ + num_recv_tokens, recv_sf_token_stride, \ + recv_sf_hidden_stride, ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + elem_size == EL && num_sf_packs == SFP && \ + num_channels == hybrid_num_channels(S)) { \ + LAUNCH_HYBRID_DISPATCH_EPILOGUE((H) * (EL), SFP, E, K, M, S, SO, SU, \ + (S) * kElasticNumHybridForwardWarps); \ + return; \ + } + +#define TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, SO, SU) \ + TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, \ + static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, 1, (H) / 128) + +#define TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE +#undef TRY_HYBRID_DISPATCH_EPILOGUE +#undef TRY_HYBRID_DISPATCH_EPILOGUE_TYPED +#undef LAUNCH_HYBRID_DISPATCH_EPILOGUE + } +#endif + +#define LAUNCH_DISPATCH_EPILOGUE(HB, SFP, E, K, M, S, R) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + auto kernel = do_expand ? \ + elastic::dispatch_copy_epilogue_impl< \ + true, false, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + (cached_mode ? \ + elastic::dispatch_copy_epilogue_impl< \ + false, true, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + elastic::dispatch_copy_epilogue_impl< \ + false, false, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K>); \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + ctx.buffer, ctx.workspace, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, recv_x, \ + static_cast(recv_sf), recv_topk_idx, \ + recv_topk_weights, recv_src_metadata, \ + channel_linked_list, num_recv_tokens, \ + recv_sf_token_stride, recv_sf_hidden_stride, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && elem_size == EL && \ + num_sf_packs == SFP) { \ + LAUNCH_DISPATCH_EPILOGUE((H) * (EL), SFP, E, K, M, S, R); \ + return; \ + } + +#define TRY_DISPATCH_EPILOGUE(H, E, K, M, S, R) \ + TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, 1, (H) / 128) + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_DISPATCH_EPILOGUE_TYPED(4096, 256, 8, 128, 24, 2, + static_cast(sizeof(nv_bfloat16)), 0); + TRY_DISPATCH_EPILOGUE_TYPED(4096, 256, 8, 128, 24, 8, + static_cast(sizeof(nv_bfloat16)), 0); +#else + TRY_DISPATCH_EPILOGUE(4096, 256, 8, 128, 24, 8); + TRY_DISPATCH_EPILOGUE(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_DISPATCH_EPILOGUE +#undef TRY_DISPATCH_EPILOGUE_TYPED +#undef LAUNCH_DISPATCH_EPILOGUE + unsupported_elastic_config("dispatch_copy_epilogue", hidden, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void* launch_mooncake_elastic_combine( + void* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + int* token_metadata_at_forward, int* channel_linked_list, + int num_reduced_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, int num_sms, int num_smem_bytes, + int num_channels, bool use_expanded_layout, bool allow_multiple_reduction, + const ElasticLaunchContext& ctx, cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, combine_smem_bytes(hidden, num_topk, kElasticNumEpilogueWarps)); + const auto comm_ctx = make_comm_ctx(ctx); + (void)token_metadata_at_forward; + (void)channel_linked_list; + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { + const int hybrid_combine_warps = + kElasticNumHybridScaleupWarps + kElasticNumHybridForwardWarps; + const int hybrid_threads = hybrid_combine_warps * 32; + const int hybrid_smem_bytes = std::max( + num_smem_bytes, + combine_smem_bytes(hidden, num_topk, hybrid_combine_warps)); + +#define LAUNCH_HYBRID_COMBINE(H, E, K, M, S, SO, SU) \ + do { \ + auto kernel = elastic::hybrid_combine_impl< \ + false, true, S, kElasticNumHybridScaleupWarps, \ + kElasticNumHybridForwardWarps, SO, SU, H, M, E, K, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, hybrid_smem_bytes, \ + stream, static_cast(x), \ + topk_weights, src_metadata, \ + psum_num_recv_tokens_per_scaleup_rank, \ + token_metadata_at_forward, channel_linked_list, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx, \ + num_reduced_tokens); \ + } while (false) + +#define TRY_HYBRID_COMBINE(H, E, K, M, S, SO, SU) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + allow_multiple_reduction && !use_expanded_layout && \ + num_channels == hybrid_num_channels(S) && \ + token_metadata_at_forward != nullptr && channel_linked_list != nullptr) { \ + LAUNCH_HYBRID_COMBINE(H, E, K, M, S, SO, SU); \ + return hybrid_combine_reduce_buffer_ptr( \ + ctx.buffer, H, K, M, SO, SU, allow_multiple_reduction); \ + } + +#define TRY_HYBRID_COMBINE_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_COMBINE(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_COMBINE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_COMBINE_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_COMBINE_SHAPE +#undef TRY_HYBRID_COMBINE +#undef LAUNCH_HYBRID_COMBINE + } +#endif + + (void)num_channels; + +#define LAUNCH_COMBINE(H, E, K, M, S, R) \ + do { \ + auto kernel = elastic::combine_impl; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + static_cast(x), topk_weights, \ + src_metadata, psum_num_recv_tokens_per_scaleup_rank,\ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.scaleup_rank_idx, num_reduced_tokens); \ + } while (false) + +#define TRY_COMBINE(H, E, K, M, S, R) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && !use_expanded_layout && \ + allow_multiple_reduction) { \ + LAUNCH_COMBINE(H, E, K, M, S, R); \ + return ctx.buffer; \ + } + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_COMBINE(4096, 256, 8, 128, 24, 2); + TRY_COMBINE(4096, 256, 8, 128, 24, 8); +#else + TRY_COMBINE(4096, 256, 8, 128, 24, 8); + TRY_COMBINE(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_COMBINE +#undef LAUNCH_COMBINE + unsupported_elastic_config("combine", hidden, num_experts, num_topk, + num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void launch_mooncake_elastic_combine_reduce_epilogue( + void* combined_x, float* combined_topk_weights, int64_t* combined_topk_idx, + int num_combined_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, void* reduce_buffer, void* bias_0, + void* bias_1, int num_sms, int num_smem_bytes, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, combine_epilogue_smem_bytes(hidden, kElasticNumEpilogueWarps)); + +#define LAUNCH_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU) \ + do { \ + auto kernel = elastic::combine_reduce_epilogue_impl< \ + false, true, S, kElasticNumEpilogueWarps, SO, SU, H, M, E, K>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + static_cast(combined_x), \ + combined_topk_weights, combined_topk_idx, \ + reduce_buffer, bias_0, bias_1, num_combined_tokens, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + !use_expanded_layout && allow_multiple_reduction) { \ + LAUNCH_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU); \ + return; \ + } + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 2); + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 8); +#else + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 8); + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 2); + +#define TRY_HYBRID_COMBINE_EPILOGUE_SHAPE(H, E, K, M, S) \ + TRY_COMBINE_EPILOGUE(H, E, K, M, S, 2, 4); \ + TRY_COMBINE_EPILOGUE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_COMBINE_EPILOGUE_SHAPE(4096, 256, 8, 128, 24); +#endif + +#undef TRY_HYBRID_COMBINE_EPILOGUE_SHAPE + +#undef TRY_COMBINE_EPILOGUE +#undef LAUNCH_COMBINE_EPILOGUE + unsupported_elastic_config("combine_reduce_epilogue", hidden, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +} // namespace mooncake diff --git a/mooncake-ep/tests/test_elastic_buffer.py b/mooncake-ep/tests/test_elastic_buffer.py new file mode 100644 index 0000000000..53f049303c --- /dev/null +++ b/mooncake-ep/tests/test_elastic_buffer.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 +"""Correctness smoke for the public Mooncake ElasticBuffer API. + +This test is intentionally self-contained: it exercises the PR's elastic +dispatch/combine path through the public ElasticBuffer wrapper and checks the +result against a deterministic PyTorch reference derived from the routing +metadata. + +Typical single-node usage: + + MOONCAKE_EP_NUM_LOCAL_RANKS=8 \ + torchrun --standalone --nproc_per_node=8 \ + mooncake-ep/tests/test_elastic_buffer.py --quick + +For multi-node hybrid validation, launch the same script with ``torchrun +--nnodes`` and set ``MOONCAKE_EP_NUM_LOCAL_RANKS`` to the number of GPUs per +node. +""" + +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.testing as testing + +from mooncake.mooncake_elastic_buffer import ElasticBuffer + + +def using_musa_backend() -> bool: + return os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in { + "1", + "ON", + "TRUE", + "YES", + } + + +def import_torchada_if_needed() -> None: + if not using_musa_backend(): + return + import torchada # noqa: F401 — maps torch.cuda.* to torch.musa.* on MUSA + + +def distributed_barrier() -> None: + if using_musa_backend(): + dist.barrier(device_ids=[torch.cuda.current_device()]) + else: + dist.barrier() + + +@dataclass(frozen=True) +class RoutePlan: + topk_idx: torch.Tensor + expected_recv_tokens: int + expected_combine_factor: int + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Test Mooncake ElasticBuffer") + parser.add_argument("--num-tokens", type=int, default=64) + parser.add_argument("--max-tokens", type=int, default=0) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--num-topk", type=int, default=8) + parser.add_argument("--num-sms", type=int, default=24) + parser.add_argument( + "--route", + choices=("alltoall", "local", "cross"), + default="alltoall", + help="Expert routing pattern to generate.", + ) + parser.add_argument( + "--allow-hybrid-mode", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument( + "--quick", + action="store_true", + help="Use smaller defaults suitable for reviewer smoke tests.", + ) + return parser.parse_args() + + +def init_distributed(seed: int) -> tuple[int, int, int]: + import_torchada_if_needed() + if not dist.is_initialized(): + dist.init_process_group("nccl") + + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count())) + torch.cuda.set_device(local_rank) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(seed + rank) + return rank, local_rank, world_size + + +def make_route_plan( + *, + rank: int, + world_size: int, + buffer: ElasticBuffer, + num_tokens: int, + num_topk: int, + num_experts: int, + route: str, +) -> RoutePlan: + local_experts = num_experts // world_size + if local_experts <= 0: + raise ValueError("num_experts must be at least world_size") + + expert_offsets = torch.arange(num_topk, device="cuda", dtype=torch.long) % local_experts + + if route == "cross" and buffer.num_scaleout_ranks > 1: + dst_scaleout = (buffer.scaleout_rank_idx + 1) % buffer.num_scaleout_ranks + dst_rank = dst_scaleout * buffer.num_scaleup_ranks + buffer.scaleup_rank_idx + choices = dst_rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + if route == "local" or (route == "cross" and buffer.num_scaleout_ranks == 1): + choices = rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + dst_ranks = (rank + torch.arange(num_topk, device="cuda", dtype=torch.long)) % world_size + choices = dst_ranks * local_experts + expert_offsets + unique_dst_ranks = int(torch.unique(dst_ranks).numel()) + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens * unique_dst_ranks, + unique_dst_ranks, + ) + + +def make_input( + *, rank: int, num_tokens: int, hidden: int, multiplier: int, addend: int = 0 +) -> torch.Tensor: + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + return (base + rank * multiplier + addend).to(torch.bfloat16).contiguous() + + +def check_dispatch_payload( + *, + rank: int, + recv_x: torch.Tensor, + handle, + expected_recv_tokens: int, + max_tokens: int, + num_tokens: int, + hidden: int, + multiplier: int, + addend: int = 0, +) -> int: + actual = int(handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if actual != expected_recv_tokens: + raise AssertionError( + f"rank={rank}: got {actual} received tokens, " + f"expected {expected_recv_tokens}" + ) + + src_global = handle.recv_src_metadata[:actual, 0].long() + src_rank = torch.div(src_global, max_tokens, rounding_mode="floor") + src_token = src_global % max_tokens + if not bool((src_token < num_tokens).all()): + raise AssertionError(f"rank={rank}: invalid source token index in metadata") + + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + expected = (base[src_token] + src_rank.view(-1, 1).float() * multiplier + addend) + expected = expected.to(torch.bfloat16) + if not torch.equal(recv_x[:actual], expected): + diff = (recv_x[:actual].float() - expected.float()).abs().max().item() + raise AssertionError(f"rank={rank}: dispatch payload mismatch, max_diff={diff}") + return actual + + +def check_combined( + *, rank: int, combined: torch.Tensor, expected: torch.Tensor, label: str +) -> None: + testing.assert_close( + combined, + expected, + rtol=5e-2, + atol=1e-3, + msg=lambda msg: f"rank={rank}: {label} combine mismatch: {msg}", + ) + + +def main() -> None: + args = parse_args() + if args.quick: + args.num_tokens = min(args.num_tokens, 32) + + rank, _local_rank, world_size = init_distributed(args.seed) + max_tokens = args.max_tokens or max(128, args.num_tokens) + num_experts = args.num_experts + + if num_experts % world_size != 0: + raise ValueError("num_experts must be divisible by world_size") + + buffer = ElasticBuffer( + dist.group.WORLD, + num_max_tokens_per_rank=max_tokens, + hidden=args.hidden, + num_topk=args.num_topk, + use_fp8_dispatch=False, + deterministic=False, + allow_hybrid_mode=args.allow_hybrid_mode, + allow_multiple_reduction=True, + num_gpu_timeout_secs=10, + ) + + route_plan = make_route_plan( + rank=rank, + world_size=world_size, + buffer=buffer, + num_tokens=args.num_tokens, + num_topk=args.num_topk, + num_experts=num_experts, + route=args.route, + ) + weights = torch.ones((args.num_tokens, args.num_topk), device="cuda", dtype=torch.float32) + + # CPU-sync dispatch: exact output extent and CPU-side expert counts. + x0 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=1_000_000, + ) + recv0, idx0, w0, handle0, _ = buffer.dispatch( + x0, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_cpu_sync=True, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + actual0 = check_dispatch_payload( + rank=rank, + recv_x=recv0, + handle=handle0, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=1_000_000, + ) + if recv0.shape[0] != actual0: + raise AssertionError(f"rank={rank}: CPU-sync dispatch returned extra rows") + if len(handle0.num_recv_tokens_per_expert_list) != num_experts // world_size: + raise AssertionError(f"rank={rank}: invalid per-expert count length") + + combined0, _, _ = buffer.combine( + recv0.contiguous(), + handle0, + topk_weights=w0[:actual0].contiguous(), + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expected0 = (x0.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined0, expected=expected0, label="sync") + + # Cached-handle dispatch: DeepEP-style path without top-k / weight inputs. + x1 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=2_000_000, + addend=17, + ) + recv1, idx1, _w1, handle1, _ = buffer.dispatch( + x1, + handle=handle0, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + actual1 = check_dispatch_payload( + rank=rank, + recv_x=recv1, + handle=handle1, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=2_000_000, + addend=17, + ) + if not torch.equal(idx1[:actual1], idx0[:actual1]): + raise AssertionError(f"rank={rank}: cached dispatch top-k metadata mismatch") + combined1, _, _ = buffer.combine( + recv1[:actual1].contiguous(), + handle1, + topk_weights=None, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expected1 = (x1.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined1, expected=expected1, label="cached") + + # Async no-CPU-sync dispatch: output keeps capacity, metadata provides exact count. + x2 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=3_000_000, + addend=31, + ) + recv2, _idx2, w2, handle2, event2 = buffer.dispatch( + x2, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_cpu_sync=False, + num_sms=args.num_sms, + async_with_compute_stream=True, + ) + event2.current_stream_wait() + torch.cuda.synchronize() + actual2 = check_dispatch_payload( + rank=rank, + recv_x=recv2, + handle=handle2, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=3_000_000, + addend=31, + ) + if handle2.num_recv_tokens_per_expert_list: + raise AssertionError(f"rank={rank}: no-CPU-sync path should not return CPU counts") + combined2, _, event3 = buffer.combine( + recv2[:actual2].contiguous(), + handle2, + topk_weights=w2[:actual2].contiguous(), + num_sms=args.num_sms, + async_with_compute_stream=True, + ) + event3.current_stream_wait() + torch.cuda.synchronize() + expected2 = (x2.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined2, expected=expected2, label="async") + + # Expanded dispatch layout: check output extent and metadata shape. + expanded_recv, expanded_idx, expanded_w, expanded_handle, _ = buffer.dispatch( + x0, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_expand=True, + do_cpu_sync=True, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expanded_actual = int(expanded_handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if expanded_actual != route_plan.expected_recv_tokens: + raise AssertionError(f"rank={rank}: expanded dispatch received-token mismatch") + expanded_output = int(expanded_handle.psum_num_recv_tokens_per_expert[-1].item()) + if expanded_recv.shape[0] != expanded_output: + raise AssertionError(f"rank={rank}: expanded output extent mismatch") + if expanded_w is not None and expanded_w.shape[0] != expanded_output: + raise AssertionError(f"rank={rank}: expanded weights extent mismatch") + if expanded_idx.shape[0] != expanded_actual: + raise AssertionError(f"rank={rank}: expanded metadata extent mismatch") + + distributed_barrier() + if rank == 0: + print( + "MOONCAKE_ELASTIC_TEST_OK", + f"world={world_size}", + f"route={args.route}", + f"recv={route_plan.expected_recv_tokens}", + f"expanded={expanded_output}", + f"scaleout={buffer.num_scaleout_ranks}", + f"scaleup={buffer.num_scaleup_ranks}", + flush=True, + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 735015a341..c7fe22818b 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -214,6 +214,7 @@ if(WITH_EP) FILES "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/ep.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_ep_buffer.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_elastic_buffer.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/pg.py" DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}) # ep.so / pg.so link against engine.so by that exact bare name. Create a diff --git a/mooncake-transfer-engine/include/transport/device/p2p_device.cuh b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh index d47ca0a2cc..824a33c2ac 100644 --- a/mooncake-transfer-engine/include/transport/device/p2p_device.cuh +++ b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh @@ -19,7 +19,11 @@ struct P2PContext { __device__ __forceinline__ bool mc_p2p_available(const P2PContext& ctx, int dst_rank) { +#ifdef MOONCAKE_EP_USE_MUSA + return ctx.peer_ptrs[dst_rank] != nullptr; +#else return ctx.available[dst_rank] != 0 && ctx.peer_ptrs[dst_rank] != nullptr; +#endif } // Translate a local pointer (within the GDR buffer) to the peer's mapped VA. diff --git a/mooncake-wheel/mooncake/mooncake_elastic_buffer.py b/mooncake-wheel/mooncake/mooncake_elastic_buffer.py new file mode 100644 index 0000000000..df751aef59 --- /dev/null +++ b/mooncake-wheel/mooncake/mooncake_elastic_buffer.py @@ -0,0 +1,582 @@ +import os +import warnings +from typing import Any, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist + +from .mooncake_ep_buffer import EventOverlap + + +def _using_musa_backend() -> bool: + return os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in { + "1", + "ON", + "TRUE", + "YES", + } + + +def _dist_barrier(group: dist.ProcessGroup) -> None: + if _using_musa_backend(): + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + else: + group.barrier() + + +def _ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def _align(x: int, alignment: int) -> int: + return _ceil_div(x, alignment) * alignment + + +class EPHandle: + """ + Official DeepEP elastic-compatible communication handle. + + The field names and semantics intentionally follow the official DeepEP elastic + handle contract so that model code can select Mooncake ElasticBuffer without + switching back to the legacy Buffer tuple handle. Mooncake stores the native + legacy handle as an implementation detail while the elastic kernels are being + wired to the Device API backend. + """ + + def __init__( + self, + do_expand: bool, + num_experts: int, + expert_alignment: int, + num_max_tokens_per_rank: int, + num_sms: int, + topk_idx: torch.Tensor, + num_recv_tokens_per_expert_list: List[int], + psum_num_recv_tokens_per_scaleup_rank: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + recv_src_metadata: torch.Tensor, + dst_buffer_slot_idx: torch.Tensor, + token_metadata_at_forward: Optional[torch.Tensor], + channel_linked_list: Optional[torch.Tensor], + native_handle: Optional[Tuple[Any, ...]] = None, + ) -> None: + assert topk_idx is not None + self.do_expand = do_expand + self.num_experts = num_experts + self.expert_alignment = expert_alignment + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_sms = num_sms + self.topk_idx = topk_idx + self.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank + self.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert + self.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list + self.recv_src_metadata = recv_src_metadata + self.dst_buffer_slot_idx = dst_buffer_slot_idx + self.token_metadata_at_forward = token_metadata_at_forward + self.channel_linked_list = channel_linked_list + self.native_handle = native_handle + + # Same convention as DeepEP: without a CPU sync this is an inferred upper + # bound; after CPU sync it tracks the actual received-token count. + self.num_recv_tokens = int(recv_src_metadata.shape[0]) + + +class ElasticBuffer: + """ + Official DeepEP elastic EP API backed by Mooncake EP transports. + + Public API source of truth: official DeepEP `ElasticBuffer`. The implementation is + deliberately separate from Mooncake's legacy `Buffer` API, while reusing the + existing Mooncake Device API transport/bootstrap path for the native data + movement backend. + """ + + # Mirrors DeepEP's fixed workspace assumptions closely enough for sizing and + # keeping one reusable buffer for all elastic EP shapes. + _NUM_MAX_RANKS = 1024 + _NUM_MAX_EXPERTS = 2048 + _NUM_MAX_CHANNELS = 8 * 160 + _NUM_BARRIER_TAGS = 16 + _NUM_MAX_INFLIGHT_AGRS = 32 + + def __init__( + self, + group: dist.ProcessGroup, + num_bytes: Optional[int] = None, + num_max_tokens_per_rank: int = 0, + hidden: int = 0, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + deterministic: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + prefer_overlap_with_compute: bool = True, + sl_idx: int = 3, + num_allocated_qps: int = 0, + num_cpu_timeout_secs: int = 300, + num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False, + ) -> None: + if not allow_multiple_reduction: + raise NotImplementedError( + "Mooncake ElasticBuffer currently supports only " + "allow_multiple_reduction=True" + ) + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.allow_hybrid_mode = allow_hybrid_mode + self.allow_multiple_reduction = allow_multiple_reduction + self.prefer_overlap_with_compute = prefer_overlap_with_compute + self.deterministic = deterministic + self.sl_idx = int(os.getenv("EP_OVERRIDE_RDMA_SL", sl_idx)) + self.num_allocated_qps = num_allocated_qps + self.num_cpu_timeout_secs = num_cpu_timeout_secs + self.num_gpu_timeout_secs = num_gpu_timeout_secs + self.explicitly_destroy = explicitly_destroy + + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.hidden = hidden + self.num_topk = num_topk + self.use_fp8_dispatch = use_fp8_dispatch + + if num_bytes is None: + num_bytes = self.get_buffer_size_hint( + group, + num_max_tokens_per_rank, + hidden, + num_topk=num_topk, + use_fp8_dispatch=use_fp8_dispatch, + allow_hybrid_mode=allow_hybrid_mode, + allow_multiple_reduction=allow_multiple_reduction, + ) + self.num_bytes = num_bytes + + ( + self.num_scaleout_ranks, + self.num_scaleup_ranks, + ) = self._calculate_logical_domain_size(group, allow_hybrid_mode) + self.scaleout_rank_idx = self.rank_idx // self.num_scaleup_ranks + self.scaleup_rank_idx = self.rank_idx % self.num_scaleup_ranks + self.num_rdma_ranks, self.num_nvlink_ranks = self._calculate_physical_domain_size(group) + + self.backend = group + + # Native Mooncake transport/runtime. This keeps the legacy Buffer ABI + # untouched while giving ElasticBuffer users a dedicated native entrypoint. + from mooncake import ep + + self.runtime = ep.ElasticBuffer( + self.rank_idx, + self.num_ranks, + num_bytes, + num_max_tokens_per_rank, + hidden, + num_topk, + use_fp8_dispatch, + deterministic, + allow_hybrid_mode, + allow_multiple_reduction, + prefer_overlap_with_compute, + self.sl_idx, + num_allocated_qps, + num_cpu_timeout_secs, + num_gpu_timeout_secs, + ) + self._connect_native() + + torch.cuda.synchronize() + _dist_barrier(group) + torch.cuda.synchronize() + + def _active_ranks_mask(self) -> list: + # `mooncake.ep.get_active_ranks` is a Mooncake PG helper and performs a + # native static cast to MooncakeBackend. ElasticBuffer transport + # bootstrap can also be driven by a regular NCCL/Gloo ProcessGroup; in + # that case every rank in the supplied group is active by definition. + if "Mooncake" not in type(self.backend).__name__: + return [1] * self.num_ranks + + from mooncake.ep import get_active_ranks + + return get_active_ranks(self.backend).tolist() + + def _connect_native(self, is_update: bool = False) -> None: + from mooncake import ep + + if not bool(self.runtime.ibgda_disabled()): + raddr, rkey = self.runtime.get_mr_info() + raddr_tensor = torch.tensor([raddr], dtype=torch.int64, device="cuda") + raddrs = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(raddrs, raddr_tensor, self.group) + raddrs_list = torch.cat(raddrs).tolist() + + rkey_tensor = torch.tensor([rkey], dtype=torch.int32, device="cuda") + rkeys = [torch.empty(1, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(rkeys, rkey_tensor, self.group) + rkeys_list = torch.cat(rkeys).tolist() + + all_to_all_size = ep.MAX_QP_COUNT // self.num_ranks + if is_update: + self.runtime.update_local_qpns() + + local_qpns = torch.tensor(self.runtime.get_local_qpns(), dtype=torch.int32, device="cuda").view( + -1, all_to_all_size + ) + remote_qpns = [torch.empty(all_to_all_size, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_to_all(remote_qpns, list(torch.unbind(local_qpns)), self.group) + peer_qpns = [remote_qpns[r].tolist() for r in range(self.num_ranks)] + + local_lids = torch.tensor(self.runtime.get_local_lids(), dtype=torch.int32, device="cuda").view( + -1, all_to_all_size + ) + remote_lids = [torch.empty(all_to_all_size, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_to_all(remote_lids, list(torch.unbind(local_lids)), self.group) + peer_lids = [remote_lids[r].tolist() for r in range(self.num_ranks)] + + subnet_prefix, interface_id = self.runtime.get_gid() + subnet_prefix_tensor = torch.tensor([subnet_prefix], dtype=torch.int64, device="cuda") + subnet_prefixes = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(subnet_prefixes, subnet_prefix_tensor, self.group) + subnet_prefixes_list = torch.cat(subnet_prefixes).tolist() + + interface_id_tensor = torch.tensor([interface_id], dtype=torch.int64, device="cuda") + interface_ids = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(interface_ids, interface_id_tensor, self.group) + interface_ids_list = torch.cat(interface_ids).tolist() + + active_ranks_mask = self._active_ranks_mask() + self.runtime.sync_ibgda_peers( + raddrs_list, + rkeys_list, + peer_qpns, + peer_lids, + subnet_prefixes_list, + interface_ids_list, + active_ranks_mask, + ) + + try: + local_handle_ints = self.runtime.get_ipc_handle() + local_handle_tensor = torch.tensor(local_handle_ints, dtype=torch.int32, device="cuda") + handles = [torch.empty(len(local_handle_ints), dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(handles, local_handle_tensor, self.group) + remote_handles = [h.tolist() for h in handles] + + active_ranks_mask = self._active_ranks_mask() + self.runtime.sync_nvlink_ipc_handles(remote_handles, active_ranks_mask) + except Exception as exc: + if bool(self.runtime.ibgda_disabled()): + raise RuntimeError( + f"[Rank {self.rank_idx}] Failed to exchange IPC handles " + "for ElasticBuffer and RDMA is disabled; native elastic " + "mode cannot continue safely." + ) from exc + warnings.warn( + f"[Rank {self.rank_idx}] Failed to exchange IPC handles for ElasticBuffer: {exc}. " + "Continuing with RDMA-only routing.", + RuntimeWarning, + stacklevel=2, + ) + + def update_ep_member(self) -> None: + self._connect_native(True) + + def destroy(self) -> None: + # Existing Mooncake Buffer owns native resources through object lifetime. + # Keep the method to match the official ElasticBuffer API. + self.runtime = None + + @staticmethod + def _workspace_num_bytes() -> int: + num_bytes = 0 + num_bytes += ElasticBuffer._NUM_BARRIER_TAGS * ( + 8 + 2 * ElasticBuffer._NUM_MAX_RANKS * 4 + ) + num_bytes += (ElasticBuffer._NUM_MAX_RANKS + ElasticBuffer._NUM_MAX_EXPERTS) * 8 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 8 * 2 + num_bytes += ElasticBuffer._NUM_MAX_EXPERTS * 8 * 2 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 4 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 4 * 2 + num_bytes += ElasticBuffer._NUM_MAX_EXPERTS * 4 * 2 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * ElasticBuffer._NUM_MAX_CHANNELS * 8 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * ElasticBuffer._NUM_MAX_CHANNELS * 4 + num_bytes += 2 * 2 * 8 + num_bytes += (ElasticBuffer._NUM_MAX_INFLIGHT_AGRS + 1) * ElasticBuffer._NUM_MAX_RANKS * 4 + return _align(num_bytes, 32) + + @staticmethod + def _atomic_scratch_num_bytes() -> int: + # Mirrors the native runtime: RDMA atomics need a local response area + # separate from the remote-visible workspace. + return ElasticBuffer._workspace_num_bytes() + + @staticmethod + def get_buffer_size_hint( + group: dist.ProcessGroup, + num_max_tokens_per_rank: int, + hidden: int, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + ) -> int: + try: + from mooncake import ep + + return int( + ep.calculate_elastic_buffer_size( + group.size(), + num_max_tokens_per_rank, + hidden, + num_topk, + use_fp8_dispatch, + allow_hybrid_mode, + allow_multiple_reduction, + ) + ) + except Exception: + pass + + num_ranks = group.size() + num_topk = max(1, num_topk) + dtype_bytes = 1 if use_fp8_dispatch else 2 + scale_bytes = _ceil_div(hidden, 128) * 4 if use_fp8_dispatch else 0 + token_bytes = _align(hidden * dtype_bytes, 32) + _align(scale_bytes, 32) + metadata_bytes = _align(num_topk * (4 + 4) + (1 + num_topk) * 4, 32) + per_slot_bytes = token_bytes + metadata_bytes + + # Direct elastic send/recv buffers plus room for combine reduce buffers. + dispatch_bytes = num_ranks * num_max_tokens_per_rank * num_topk * per_slot_bytes * 2 + combine_factor = 3 if allow_multiple_reduction else 4 + combine_bytes = dispatch_bytes * combine_factor + hybrid_factor = 2 if allow_hybrid_mode and num_ranks > 1 else 1 + return int( + ElasticBuffer._workspace_num_bytes() + + ElasticBuffer._atomic_scratch_num_bytes() + + hybrid_factor * (dispatch_bytes + combine_bytes) + ) + + @staticmethod + def get_engram_storage_size_hint( + num_entries: int, + hidden: int, + num_max_tokens_per_rank: int, + dtype: torch.dtype = torch.bfloat16, + ) -> int: + num_sf_packs = _ceil_div(hidden, 128) if dtype.itemsize <= 1 else 0 + num_bytes_per_entry = _align(hidden * dtype.itemsize + num_sf_packs * 4, 32) + return num_bytes_per_entry * (num_entries + num_max_tokens_per_rank) + + @staticmethod + def get_pp_buffer_size_hint(num_max_tensor_bytes: int, num_max_inflight_tensors: int) -> int: + return _align(num_max_tensor_bytes, 32) * num_max_inflight_tensors * 2 * 2 + + @staticmethod + def get_agrs_buffer_size_hint(group: dist.ProcessGroup, num_max_session_bytes: int) -> int: + return num_max_session_bytes + + @staticmethod + def _calculate_physical_domain_size(group: dist.ProcessGroup) -> Tuple[int, int]: + num_ranks = group.size() + num_local_ranks = int(os.getenv("MOONCAKE_EP_NUM_LOCAL_RANKS", "0")) + if num_local_ranks <= 0: + try: + num_local_ranks = max(1, min(num_ranks, torch.cuda.device_count())) + except Exception: + num_local_ranks = 1 + num_local_ranks = max(1, min(num_local_ranks, num_ranks)) + return _ceil_div(num_ranks, num_local_ranks), num_local_ranks + + @staticmethod + def _calculate_logical_domain_size(group: dist.ProcessGroup, allow_hybrid_mode: bool = True) -> Tuple[int, int]: + num_ranks = group.size() + num_rdma_ranks, num_nvlink_ranks = ElasticBuffer._calculate_physical_domain_size(group) + if allow_hybrid_mode and num_rdma_ranks > 1: + return num_rdma_ranks, num_nvlink_ranks + return 1, num_ranks + + def get_physical_domain_size(self) -> Tuple[int, int]: + return self.num_rdma_ranks, self.num_nvlink_ranks + + def get_logical_domain_size(self) -> Tuple[int, int]: + return self.num_scaleout_ranks, self.num_scaleup_ranks + + def barrier(self, use_comm_stream: bool = True, with_cpu_sync: bool = False) -> None: + if with_cpu_sync: + torch.cuda.synchronize() + _dist_barrier(self.group) + if with_cpu_sync: + torch.cuda.synchronize() + + @staticmethod + def capture() -> Any: + from mooncake import ep + + return ep.EventHandle() + + def get_theoretical_num_sms(self, num_experts: int, num_topk: int) -> int: + device = torch.cuda.current_device() + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + if self.prefer_overlap_with_compute: + return max(1, min(24, sm_count // 4)) + return max(1, min(40, sm_count // 2, num_experts * max(1, num_topk))) + + def dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + num_max_tokens_per_rank: Optional[int] = None, + expert_alignment: Optional[int] = None, + handle: Optional[EPHandle] = None, + do_expand: bool = False, + do_cpu_sync: Optional[bool] = None, + num_sms: Optional[int] = None, + async_with_compute_stream: bool = False, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], EPHandle, EventOverlap]: + if self.runtime is None: + raise RuntimeError("ElasticBuffer has been destroyed") + if handle is not None: + if topk_idx is not None or topk_weights is not None: + raise AssertionError("topk_idx and topk_weights must be None when cached handle is provided") + if do_cpu_sync: + raise AssertionError("Cannot do CPU sync with cached handle") + if handle.native_handle is None: + raise RuntimeError("Cached EPHandle is missing its native Mooncake handle") + topk_idx = handle.topk_idx + num_max_tokens_per_rank = num_max_tokens_per_rank or handle.num_max_tokens_per_rank + num_experts = num_experts or handle.num_experts + expert_alignment = handle.expert_alignment if expert_alignment is None else expert_alignment + num_sms = handle.num_sms if num_sms is None else num_sms + do_cpu_sync = False + else: + if topk_idx is None: + raise AssertionError("topk_idx must be provided when cached handle is not provided") + expert_alignment = 1 if expert_alignment is None else expert_alignment + do_cpu_sync = True if do_cpu_sync is None else do_cpu_sync + if do_expand: + warnings.warn( + "do_expand=True was requested. Mooncake currently returns the native packed expert layout; " + "expanded contiguous expert layout will be produced by the native elastic kernels.", + RuntimeWarning, + stacklevel=2, + ) + + x_data = x[0] if isinstance(x, tuple) else x + sf = x[1] if isinstance(x, tuple) else None + if num_experts is None: + num_experts = int(torch.max(topk_idx).item()) + 1 + if num_max_tokens_per_rank is None: + num_max_tokens_per_rank = self.num_max_tokens_per_rank or x_data.shape[0] + if num_sms is None: + num_sms = self.get_theoretical_num_sms(num_experts, topk_idx.shape[1]) + + active_ranks = torch.ones(self.num_ranks, dtype=torch.int32, device=x_data.device) + output = self.runtime.dispatch( + x_data, + sf, + topk_idx, + topk_weights, + active_ranks, + num_experts, + num_max_tokens_per_rank, + expert_alignment, + num_sms, + do_expand, + do_cpu_sync, + async_with_compute_stream, + handle.native_handle if handle is not None else None, + ) + native_handle = output.handle + + elastic_handle = EPHandle( + do_expand=native_handle.do_expand, + num_experts=native_handle.num_experts, + expert_alignment=native_handle.expert_alignment, + num_max_tokens_per_rank=native_handle.num_max_tokens_per_rank, + num_sms=native_handle.num_sms, + topk_idx=native_handle.topk_idx, + num_recv_tokens_per_expert_list=list(native_handle.num_recv_tokens_per_expert_list), + psum_num_recv_tokens_per_scaleup_rank=native_handle.psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert=native_handle.psum_num_recv_tokens_per_expert, + recv_src_metadata=native_handle.recv_src_metadata, + dst_buffer_slot_idx=native_handle.dst_buffer_slot_idx, + token_metadata_at_forward=native_handle.token_metadata_at_forward, + channel_linked_list=native_handle.channel_linked_list, + native_handle=native_handle, + ) + recv_x = (output.recv_x, output.recv_x_scales) if output.recv_x_scales is not None else output.recv_x + tensors_to_record = ( + x_data, + topk_idx, + active_ranks, + output.recv_x, + output.recv_topk_idx, + native_handle.topk_idx, + native_handle.psum_num_recv_tokens_per_scaleup_rank, + native_handle.psum_num_recv_tokens_per_expert, + native_handle.recv_src_metadata, + native_handle.dst_buffer_slot_idx, + *(() if sf is None else (sf,)), + *(() if topk_weights is None else (topk_weights,)), + *(() if output.recv_x_scales is None else (output.recv_x_scales,)), + *(() if output.recv_topk_weights is None else (output.recv_topk_weights,)), + *(() if native_handle.token_metadata_at_forward is None else (native_handle.token_metadata_at_forward,)), + *(() if native_handle.channel_linked_list is None else (native_handle.channel_linked_list,)), + ) + return ( + recv_x, + output.recv_topk_idx, + output.recv_topk_weights, + elastic_handle, + EventOverlap(output.event, tensors_to_record if async_with_compute_stream else None), + ) + + def combine( + self, + x: torch.Tensor, + handle: EPHandle, + topk_weights: Optional[torch.Tensor] = None, + num_sms: Optional[int] = None, + async_with_compute_stream: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]: + if self.runtime is None: + raise RuntimeError("ElasticBuffer has been destroyed") + if handle.native_handle is None: + raise RuntimeError("Mooncake EPHandle does not contain a native handle") + active_ranks = torch.ones(self.num_ranks, dtype=torch.int32, device=x.device) + if topk_weights is None: + topk_weights = torch.ones_like(handle.topk_idx, dtype=torch.float32, device=x.device) + output = self.runtime.combine( + x, + handle.native_handle, + topk_weights, + active_ranks, + num_sms if num_sms is not None else handle.num_sms, + async_with_compute_stream, + None, + ) + native_handle = handle.native_handle + tensors_to_record = ( + x, + topk_weights, + active_ranks, + output.combined_x, + native_handle.topk_idx, + native_handle.psum_num_recv_tokens_per_scaleup_rank, + native_handle.psum_num_recv_tokens_per_expert, + native_handle.recv_src_metadata, + native_handle.dst_buffer_slot_idx, + *(() if native_handle.token_metadata_at_forward is None else (native_handle.token_metadata_at_forward,)), + *(() if native_handle.channel_linked_list is None else (native_handle.channel_linked_list,)), + ) + return ( + output.combined_x, + output.combined_topk_weights, + EventOverlap(output.event, tensors_to_record if async_with_compute_stream else None), + ) + + +__all__ = ["ElasticBuffer", "EPHandle", "EventOverlap"] From 7192ff4fbe5c5015e060080b24063677c4f833bc Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Mon, 13 Jul 2026 09:58:18 +0800 Subject: [PATCH 55/57] [TE] Add host control fallback for IBGDA QP setup (#2867) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../include/transport/device/ibgda/mlx5gda.h | 10 ++ .../device/ibgda_device_transport.cpp | 144 ++++++++++++++---- .../src/transport/device/mlx5gda.cpp | 96 ++++++++++-- 3 files changed, 202 insertions(+), 48 deletions(-) diff --git a/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h index 82030954d2..34584ff3ad 100644 --- a/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h +++ b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h @@ -87,6 +87,16 @@ struct mlx5gda_qp_devctx { uint16_t wq_tail; // last non-completed wqeid }; +struct mlx5gda_create_qp_failure { + bool valid; + uint32_t status; + uint32_t syndrome; + int sys_errno; +}; + +void mlx5gda_reset_create_qp_failure(); +mlx5gda_create_qp_failure mlx5gda_last_create_qp_failure(); + struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf, struct mlx5dv_devx_umem *ctrl_buf_umem, struct memheap *ctrl_buf_heap, diff --git a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp index 2dcdac17dc..bc8e20dbd5 100644 --- a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp +++ b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp @@ -25,11 +25,13 @@ #include #include +#include #include #include #include "cuda_alike.h" #include "transport/device/ibgda/memheap.h" +#include "transport/device/ibgda/mlx5_ifc.h" #include "transport/device/ibgda/mlx5gda.h" #include "topology.h" @@ -199,28 +201,7 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { } int allocateControlBuffer() override { - cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize); - if (err != cudaSuccess) { - LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: " - << cudaGetErrorString(err); - return -1; - } - - ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize, - IBV_ACCESS_LOCAL_WRITE); - if (!ctrl_buf_umem_) { - LOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg failed (errno=" - << errno << ")"; - return -1; - } - LOG(INFO) << "[EP IBGDA] ctrl_buf UMEM registered via VA path"; - - ctrl_buf_heap_ = memheap_create(kCtrlBufSize); - if (!ctrl_buf_heap_) { - LOG(ERROR) << "[EP IBGDA] memheap_create failed"; - return -1; - } - return 0; + return allocateControlBuffer(false); } int createQueuePairs(void* stream_ptr) override { @@ -231,6 +212,8 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { ctrl_buf_heap_, pd_, 16384, 1, stream); if (!qp) { LOG(ERROR) << "[EP IBGDA] mlx5gda_create_rc_qp failed at " << i; + if (retryWithHostControlBuffer()) + return createQueuePairs(stream_ptr); return -1; } if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) { @@ -243,11 +226,11 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { .qpn = qp->qpn, .wqeid_mask = qp->num_wqebb - 1, .wq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->wq_offset), + static_cast(ctrl_buf_dev_) + qp->wq_offset), .cq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->send_cq->cq_offset), + static_cast(ctrl_buf_dev_) + qp->send_cq->cq_offset), .dbr = reinterpret_cast( - static_cast(ctrl_buf_) + qp->dbr_offset), + static_cast(ctrl_buf_dev_) + qp->dbr_offset), .bf = static_cast(qp->uar->reg_addr), }; cudaMemcpy( @@ -259,11 +242,7 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { } int recreateQueuePairs(void* stream_ptr) override { - auto stream = static_cast(stream_ptr); - for (auto* qp : qps_) { - if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp); - } - qps_.clear(); + destroyQueuePairs(); return createQueuePairs(stream_ptr); } @@ -349,11 +328,98 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { int gidIndex() const override { return gid_index_; } private: - void teardown() { + static bool isCreateQpBadParam(const mlx5gda_create_qp_failure& failure) { + return failure.valid && failure.status == MLX5_CMD_STAT_BAD_PARAM_ERR; + } + + int allocateControlBuffer(bool host_backed) { + ctrl_buf_host_ = host_backed; + if (ctrl_buf_host_) { + void* ptr = nullptr; + int ret = posix_memalign(&ptr, 4096, kCtrlBufSize); + if (ret != 0) { + LOG(ERROR) << "[EP IBGDA] posix_memalign ctrl_buf failed: " + << ret; + return -1; + } + ctrl_buf_ = ptr; + std::memset(ctrl_buf_, 0, kCtrlBufSize); + + cudaError_t err = cudaHostRegister( + ctrl_buf_, kCtrlBufSize, + cudaHostRegisterPortable | cudaHostRegisterMapped); + if (err != cudaSuccess) { + LOG(ERROR) << "[EP IBGDA] cudaHostRegister ctrl_buf failed: " + << cudaGetErrorString(err); + free(ctrl_buf_); + ctrl_buf_ = nullptr; + return -1; + } + + err = cudaHostGetDevicePointer(&ctrl_buf_dev_, ctrl_buf_, 0); + if (err != cudaSuccess) { + LOG(ERROR) + << "[EP IBGDA] cudaHostGetDevicePointer ctrl_buf failed: " + << cudaGetErrorString(err); + cudaHostUnregister(ctrl_buf_); + free(ctrl_buf_); + ctrl_buf_ = nullptr; + ctrl_buf_host_ = false; + return -1; + } + LOG(INFO) << "[EP IBGDA] Using host-backed mapped control buffer"; + } else { + cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize); + if (err != cudaSuccess) { + LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: " + << cudaGetErrorString(err); + return -1; + } + ctrl_buf_dev_ = ctrl_buf_; + } + + ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize, + IBV_ACCESS_LOCAL_WRITE); + if (!ctrl_buf_umem_) { + LOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg failed (errno=" + << errno << ")"; + freeControlBuffer(); + return -1; + } + LOG(INFO) << "[EP IBGDA] ctrl_buf UMEM registered via VA path"; + + ctrl_buf_heap_ = memheap_create(kCtrlBufSize); + if (!ctrl_buf_heap_) { + LOG(ERROR) << "[EP IBGDA] memheap_create failed"; + freeControlBuffer(); + return -1; + } + return 0; + } + + bool retryWithHostControlBuffer() { + auto failure = mlx5gda_last_create_qp_failure(); + if (ctrl_buf_host_ || !isCreateQpBadParam(failure)) return false; + + LOG(WARNING) << "[EP IBGDA] GPU-backed control buffer was rejected by " + "DevX CREATE_QP" + << " (status=0x" << std::hex << failure.status + << " syndrome=0x" << failure.syndrome << std::dec + << "); retrying with host-backed mapped control buffer"; + + destroyQueuePairs(); + freeControlBuffer(); + return allocateControlBuffer(true) == 0; + } + + void destroyQueuePairs() { for (auto* qp : qps_) { if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp); } qps_.clear(); + } + + void freeControlBuffer() { if (ctrl_buf_heap_) { memheap_destroy(ctrl_buf_heap_); ctrl_buf_heap_ = nullptr; @@ -363,9 +429,21 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { ctrl_buf_umem_ = nullptr; } if (ctrl_buf_) { - cudaFree(ctrl_buf_); + if (ctrl_buf_host_) { + cudaHostUnregister(ctrl_buf_); + free(ctrl_buf_); + } else { + cudaFree(ctrl_buf_); + } ctrl_buf_ = nullptr; + ctrl_buf_dev_ = nullptr; + ctrl_buf_host_ = false; } + } + + void teardown() { + destroyQueuePairs(); + freeControlBuffer(); if (mr_) { ibv_dereg_mr(mr_); mr_ = nullptr; @@ -407,6 +485,8 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { // Control buffer void* ctrl_buf_ = nullptr; // GPU VA + void* ctrl_buf_dev_ = nullptr; + bool ctrl_buf_host_ = false; mlx5dv_devx_umem* ctrl_buf_umem_ = nullptr; memheap* ctrl_buf_heap_ = nullptr; diff --git a/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp index ca601b2536..532e627027 100644 --- a/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp +++ b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -38,6 +39,47 @@ static void print_cuda_error(const char* msg) { fprintf(stderr, "%s: %s\n", msg, err_str); } +static thread_local mlx5gda_create_qp_failure g_last_create_qp_failure{}; + +void mlx5gda_reset_create_qp_failure() { g_last_create_qp_failure = {}; } + +mlx5gda_create_qp_failure mlx5gda_last_create_qp_failure() { + return g_last_create_qp_failure; +} + +static bool is_host_control_buffer(void* ptr) { + cudaPointerAttributes attr{}; + cudaError_t err = cudaPointerGetAttributes(&attr, ptr); + if (err != cudaSuccess) { + cudaGetLastError(); + return false; + } + return attr.type == cudaMemoryTypeHost; +} + +static void print_devx_create_qp_failure(const uint8_t* cmd_out, + uint32_t num_wqebb, uint8_t port_num, + uint32_t pdn, uint32_t uar_page, + uint32_t cqn, uint32_t umem_id, + size_t wq_offset, size_t dbr_offset, + int sys_errno) { + uint32_t status = DEVX_GET(create_qp_out, cmd_out, status); + uint32_t syndrome = DEVX_GET(create_qp_out, cmd_out, syndrome); + g_last_create_qp_failure = mlx5gda_create_qp_failure{ + .valid = true, + .status = status, + .syndrome = syndrome, + .sys_errno = sys_errno, + }; + fprintf(stderr, + "mlx5dv_devx_obj_create(create_qp) failed: errno=%d (%s) " + "status=0x%x syndrome=0x%x port=%u pdn=0x%x uar_page=0x%x " + "cqn=0x%x umem_id=0x%x num_wqebb=%u wq_offset=0x%zx " + "dbr_offset=0x%zx\n", + sys_errno, strerror(sys_errno), status, syndrome, port_num, pdn, + uar_page, cqn, umem_id, num_wqebb, wq_offset, dbr_offset); +} + // Create UAR for BF (Blue Flame) doorbell ringing. // On CUDA: registers the BF MMIO region into GPU address space so the // GPU kernel can directly write the doorbell (lowest latency). @@ -99,6 +141,7 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, struct ibv_context* ctx = pd->context; void* cq_context = NULL; + bool ctrl_host = is_host_control_buffer(ctrl_buf); if (cqe <= 0) { errno = EINVAL; @@ -120,11 +163,16 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, // initialized to 0xFF (-1) to mark them as invalid. The hardware checks the // owner bit in CQE to determine if it's valid. This is mandatory for proper // CQ operation. Use async version to avoid blocking. - if (cudaMemsetAsync(ctrl_buf + cq_offset, -1, - num_cqe * sizeof(struct mlx5_cqe64), - stream) != cudaSuccess) { - print_cuda_error("Failed to memset CQ memory"); - goto fail; + if (ctrl_host) { + memset(static_cast(ctrl_buf) + cq_offset, -1, + num_cqe * sizeof(struct mlx5_cqe64)); + } else { + if (cudaMemsetAsync(static_cast(ctrl_buf) + cq_offset, -1, + num_cqe * sizeof(struct mlx5_cqe64), + stream) != cudaSuccess) { + print_cuda_error("Failed to memset CQ memory"); + goto fail; + } } dbr_offset = memheap_alloc(ctrl_buf_heap, sizeof(struct mlx5gda_cq_dbr)); if (dbr_offset == -1) { @@ -162,9 +210,11 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, // Synchronize stream before creating CQ object, as hardware will read the // CQE memory - if (cudaStreamSynchronize(stream) != cudaSuccess) { - print_cuda_error("Failed to synchronize stream before CQ creation"); - goto fail; + if (!ctrl_host) { + if (cudaStreamSynchronize(stream) != cudaSuccess) { + print_cuda_error("Failed to synchronize stream before CQ creation"); + goto fail; + } } mlx5_cq = mlx5dv_devx_obj_create(ctx, cmd_in, sizeof(cmd_in), cmd_out, @@ -208,6 +258,8 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, struct memheap* ctrl_buf_heap, struct ibv_pd* pd, int wqe, uint8_t port_num, cudaStream_t stream) { + mlx5gda_reset_create_qp_failure(); + struct mlx5gda_qp* qp = NULL; struct mlx5gda_cq* send_cq = NULL; struct mlx5dv_devx_uar* uar = NULL; @@ -219,6 +271,7 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, void* qp_context = NULL; void* cap = NULL; uint32_t cqe_version = 0; + bool ctrl_host = is_host_control_buffer(ctrl_buf); if (wqe <= 0) { errno = EINVAL; @@ -291,10 +344,16 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, goto fail; } // DBR must be zero-initialized. Use async version to avoid blocking. - if (cudaMemsetAsync(ctrl_buf + dbr_offset, 0, sizeof(struct mlx5gda_wq_dbr), - stream) != cudaSuccess) { - print_cuda_error("Failed to zero DBR memory"); - goto fail; + if (ctrl_host) { + memset(static_cast(ctrl_buf) + dbr_offset, 0, + sizeof(struct mlx5gda_wq_dbr)); + } else { + if (cudaMemsetAsync(static_cast(ctrl_buf) + dbr_offset, 0, + sizeof(struct mlx5gda_wq_dbr), + stream) != cudaSuccess) { + print_cuda_error("Failed to zero DBR memory"); + goto fail; + } } DEVX_SET(create_qp_in, cmd_in, opcode, MLX5_CMD_OP_CREATE_QP); @@ -326,14 +385,19 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, // Synchronize stream before creating QP object, as hardware will read the // DBR memory - if (cudaStreamSynchronize(stream) != cudaSuccess) { - print_cuda_error("Failed to synchronize stream before QP creation"); - goto fail; + if (!ctrl_host) { + if (cudaStreamSynchronize(stream) != cudaSuccess) { + print_cuda_error("Failed to synchronize stream before QP creation"); + goto fail; + } } mlx5_qp = mlx5dv_devx_obj_create(ctx, cmd_in, sizeof(cmd_in), cmd_out, sizeof(cmd_out)); if (mlx5_qp == NULL) { + print_devx_create_qp_failure( + cmd_out, num_wqebb, port_num, mpd.pdn, uar->page_id, send_cq->cqn, + ctrl_buf_umem->umem_id, wq_offset, dbr_offset, errno); goto fail; } @@ -517,4 +581,4 @@ int mlx5gda_modify_rc_qp_rtr2rts(struct mlx5gda_qp* qp) { perror("Failed to modify RC QP (rtr2rts)"); } return ret; -} \ No newline at end of file +} From 60e6d338d23983a42564e7d55a02340211ae99ed Mon Sep 17 00:00:00 2001 From: liangxu Date: Thu, 16 Jul 2026 16:32:18 +0800 Subject: [PATCH 56/57] feat: LOCAL_DISK drain migration via Move flow (Option A) Enable drain jobs to migrate LOCAL_DISK replicas by reusing the existing Move flow with a new LOCAL_DISK transfer path: - replica.h: get_segment_names() returns transport_endpoint for LOCAL_DISK - client_service: ExecuteReplicaTransfer adds LOCAL_DISK branch (SSD -> staging buffer -> TransferWrite -> target MEMORY) - file_storage: expose LoadBatchFromLocalDisk (AllocateBatch + BatchLoad) for client access - master_service: CreateMoveTask falls back to LOCAL_DISK client lookup; ScheduleDrainJobTasks scans LOCAL_DISK replicas with MEMORY-first dedup; MaybeCompleteDrainJob checks LOCAL_DISK residuals Reference pattern: ProcessPromotionTasks. No new RPCs, no new master state. --- mooncake-store/include/client_service.h | 19 ++- mooncake-store/include/file_storage.h | 28 ++++- mooncake-store/include/master_service.h | 20 +++ mooncake-store/include/replica.h | 6 + mooncake-store/src/client_service.cpp | 83 ++++++++++++- mooncake-store/src/file_storage.cpp | 40 ++++++ mooncake-store/src/master_service.cpp | 159 +++++++++++++++++++++--- 7 files changed, 332 insertions(+), 23 deletions(-) diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 621b4edd70..156b9c8129 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -28,6 +28,11 @@ #include "local_hot_cache.h" #include "pinned_buffer_pool.h" +namespace mooncake { +class FileStorage; +} + + namespace mooncake { class PutOperation; @@ -67,6 +72,13 @@ class Client { const UUID& getClientId() const { return client_id_; } const std::string& tenant_id() const { return master_client_.tenant_id(); } + // Set by FileStorage during construction so that Client can access + // local-disk loading facilities (AllocateBatch + BatchLoad) for + // LOCAL_DISK drain migration. Raw pointer: FileStorage owns a + // shared_ptr, so the lifetime is safe as long as + // FileStorage is alive. + void set_file_storage(FileStorage* fs) { file_storage_ = fs; } + /** * @brief Creates and initializes a new Client instance * @param local_hostname Local host address (IP:Port) @@ -829,6 +841,10 @@ class Client { ThreadPool write_thread_pool_; std::shared_ptr storage_backend_; + // Non-owning pointer to FileStorage, set via set_file_storage(). + // Used for LOCAL_DISK drain migration to access AllocateBatch/BatchLoad. + FileStorage* file_storage_ = nullptr; + // For high availability std::unique_ptr leader_coordinator_; std::mutex leader_switch_mutex_; @@ -863,7 +879,8 @@ class Client { void ExecuteTask(const ClientTask& client_task); tl::expected ExecuteReplicaTransfer( - const std::string& key, const std::string& action_name, + const std::string& key, const std::string& tenant_id, + const std::string& action_name, std::function()> end_fn, std::function()> revoke_fn, const Replica::Descriptor& source, diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 4a2f996310..a0ccb71ddf 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -49,9 +49,11 @@ class FileStorage { */ bool ReleaseBuffer(uint64_t batch_id); - private: - friend class FileStorageTest; - friend class FileStoragePromotionTest; + /** + * @brief RAII wrapper for a batch of O_DIRECT-aligned staging buffers + * allocated from client_buffer_allocator_. The BufferHandles in + * `handles` release the staging space when this object is destroyed. + */ struct AllocatedBatch { uint64_t batch_id; std::vector handles; @@ -70,6 +72,26 @@ class FileStorage { ~AllocatedBatch() = default; }; + /** + * @brief Load a single key from local SSD into a staging buffer. + * Used by Client::ExecuteReplicaTransfer for LOCAL_DISK drain + * migration. Allocates an O_DIRECT-aligned buffer via + * client_buffer_allocator_, reads the data from the local SSD + * backend, and returns the batch. The caller is responsible for + * keeping the shared_ptr alive while using the slices (RAII + * releases the staging space when it goes out of scope). + * + * @param key Object key (without tenant prefix) + * @param tenant_id Tenant ID (used to build the storage key) + * @param size Expected data size in bytes + * @return shared_ptr on success, error on failure + */ + tl::expected, ErrorCode> + LoadBatchFromLocalDisk(const std::string& key, + const std::string& tenant_id, uint64_t size); + + private: + /** * @brief Offload object data and metadata. * @return tl::expected indicating operation status. diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 42fd27efa3..e79f22ef84 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -1025,6 +1025,13 @@ class MasterService { return it != replicas_.end() ? &(*it) : nullptr; } + const Replica* GetFirstReplica( + const std::function& pred_fn) const { + const auto it = + std::find_if(replicas_.begin(), replicas_.end(), pred_fn); + return it != replicas_.end() ? &(*it) : nullptr; + } + Replica* GetReplicaByID(const ReplicaID& id) { return GetFirstReplica( [&id](const Replica& replica) { return replica.id() == id; }); @@ -1074,6 +1081,19 @@ class MasterService { }); } + const Replica* GetReplicaBySegmentName( + const std::string& segment_name) const { + return GetFirstReplica([&segment_name](const Replica& replica) { + auto names = replica.get_segment_names(); + for (auto& name_opt : names) { + if (name_opt == segment_name) { + return true; + } + } + return false; + }); + } + // Grant a lease with timeout as now() + ttl, only update if the new // timeout is larger void GrantLease(const uint64_t ttl, const uint64_t soft_ttl) const { diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index c3acf46dcf..0dd451f67a 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -653,6 +653,12 @@ inline std::vector> Replica::get_segment_names() segment_names.push_back(std::nullopt); } return segment_names; + } else if (is_local_disk_replica()) { + // LOCAL_DISK replicas use transport_endpoint as their segment + // identifier so that GetReplicaBySegmentName can locate them for + // drain-driven Move operations. + const auto& disk_data = std::get(data_); + return {disk_data.transport_endpoint}; } return std::vector>(); } diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 446503c91a..55189ce617 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -3,6 +3,7 @@ #include #include "allocator.h" +#include "file_storage.h" #include "segment.h" #include @@ -3164,7 +3165,8 @@ tl::expected Client::CreateMoveTask( } tl::expected Client::ExecuteReplicaTransfer( - const std::string& key, const std::string& action_name, + const std::string& key, const std::string& tenant_id, + const std::string& action_name, std::function()> end_fn, std::function()> revoke_fn, const Replica::Descriptor& source, @@ -3178,7 +3180,78 @@ tl::expected Client::ExecuteReplicaTransfer( } }; - // currently only memory source replica is supported + // === LOCAL_DISK source path === + // Load data from local SSD into a staging buffer, then TransferWrite + // to each target. This mirrors the ProcessPromotionTasks pattern: + // AllocateBatch (O_DIRECT-aligned) -> BatchLoad -> TransferWrite. + if (source.is_local_disk_replica()) { + if (file_storage_ == nullptr) { + LOG(ERROR) << "action=replica_" << action_name << "_failed" + << ", key=" << key + << ", error=no_file_storage_for_local_disk_source"; + revoke_lambda(); + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + const auto& local_disk_desc = source.get_local_disk_descriptor(); + const uint64_t object_size = local_disk_desc.object_size; + + // (a) Load from local SSD into a staging buffer. + auto load_result = + file_storage_->LoadBatchFromLocalDisk(key, tenant_id, object_size); + if (!load_result) { + LOG(ERROR) << "action=replica_" << action_name << "_failed" + << ", key=" << key + << ", error=local_disk_load_failed" + << ", error_code=" << load_result.error(); + revoke_lambda(); + return tl::unexpected(load_result.error()); + } + + // Keep the batch alive for the duration of the transfer. + // The shared_ptr RAII-releases the staging buffer when it goes + // out of scope. + auto staging = std::move(load_result.value()); + + // The storage key used by FileStorage is tenant-scoped. Find the + // slice that corresponds to our key. + const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key); + auto slice_it = staging->slices.find(storage_key); + if (slice_it == staging->slices.end()) { + LOG(ERROR) << "action=replica_" << action_name << "_failed" + << ", key=" << key + << ", error=staging_slice_missing"; + revoke_lambda(); + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + + // (b) Split the staging buffer into transfer slices. + auto slices = split_into_slices(slice_it->second.ptr, + slice_it->second.size); + + // (c) TransferWrite to each target. + for (const auto& target : targets) { + if (TransferWrite(target, slices) != ErrorCode::OK) { + LOG(ERROR) << "action=replica_" << action_name << "_failed" + << ", key=" << key + << ", error=transfer_write_failed"; + revoke_lambda(); + return tl::unexpected(ErrorCode::TRANSFER_FAIL); + } + } + + // (d) Finalize. The staging buffer is released when staging goes + // out of scope at the end of this block. + auto end_result = end_fn(); + if (!end_result.has_value()) { + revoke_lambda(); + return tl::unexpected(end_result.error()); + } + + return {}; + } + + // === MEMORY source path (existing logic) === if (!source.is_memory_replica()) { LOG(ERROR) << "action=replica_" << action_name << "_failed" << ", key=" << key << ", error=invalid_replica_type"; @@ -3259,7 +3332,8 @@ tl::expected Client::Copy( } auto result = ExecuteReplicaTransfer( - key, "copy", [&]() { return master_client_.CopyEnd(key, tenant_id); }, + key, tenant_id, "copy", + [&]() { return master_client_.CopyEnd(key, tenant_id); }, [&]() { return master_client_.CopyRevoke(key, tenant_id); }, response.source, response.targets); @@ -3314,7 +3388,8 @@ tl::expected Client::Move(const std::string& key, std::vector targets = {response.target.value()}; auto result = ExecuteReplicaTransfer( - key, "move", [&]() { return master_client_.MoveEnd(key, tenant_id); }, + key, tenant_id, "move", + [&]() { return master_client_.MoveEnd(key, tenant_id); }, [&]() { return master_client_.MoveRevoke(key, tenant_id); }, response.source, targets); diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 05d7ef2e02..999137c412 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -214,6 +214,13 @@ FileStorage::FileStorage(const FileStorageConfig& config, } } #endif + + // Register this FileStorage with the Client so that + // ExecuteReplicaTransfer can access local-disk loading for + // LOCAL_DISK drain migration. + if (client_) { + client_->set_file_storage(this); + } } FileStorage::~FileStorage() { @@ -1019,6 +1026,39 @@ FileStorage::AllocateBatch(const std::vector& keys, return result; } +tl::expected, ErrorCode> +FileStorage::LoadBatchFromLocalDisk(const std::string& key, + const std::string& tenant_id, + uint64_t size) { + if (size == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key); + + // (a) Allocate an O_DIRECT-aligned staging buffer. + std::vector single_key{storage_key}; + std::vector single_size{static_cast(size)}; + auto allocate_res = AllocateBatch(single_key, single_size); + if (!allocate_res) { + LOG(WARNING) << "LoadBatchFromLocalDisk: AllocateBatch failed for key=" + << key << ", error=" << allocate_res.error(); + return tl::make_unexpected(allocate_res.error()); + } + + auto staging = allocate_res.value(); + + // (b) Read the data from local SSD into the staging buffer. + auto load_res = BatchLoad(staging->slices); + if (!load_res) { + LOG(WARNING) << "LoadBatchFromLocalDisk: BatchLoad failed for key=" + << key << ", error=" << load_res.error(); + return tl::make_unexpected(load_res.error()); + } + + return staging; +} + void FileStorage::ClientBufferGCThreadFunc() { LOG(INFO) << "action=client_buffer_gc_thread_started"; while (client_buffer_gc_running_) { diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index bf55656980..d1d1ec3944 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -8221,6 +8221,22 @@ tl::expected MasterService::CreateMoveTask( ErrorCode error = segment_accessor.GetClientIdBySegmentName(source, select_client); + // If segment-name lookup failed, the source may be a LOCAL_DISK + // transport_endpoint (returned by Replica::get_segment_names() for + // LOCAL_DISK replicas). Fall back to finding the client_id from + // the LOCAL_DISK replica itself. + if (error != ErrorCode::OK) { + auto local_disk_replica = metadata.GetReplicaBySegmentName(source); + if (local_disk_replica && local_disk_replica->is_local_disk_replica()) { + auto client_id_opt = + local_disk_replica->get_local_disk_client_id(); + if (client_id_opt.has_value()) { + select_client = client_id_opt.value(); + error = ErrorCode::OK; + } + } + } + if (error != ErrorCode::OK) { LOG(ERROR) << "key=" << key << ", segment_name=" << source << ", error=client_id_not_found"; @@ -8555,37 +8571,103 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { } std::unordered_set blocked_unit_keys; + + // Pre-compute drain segment -> client_id mapping. Used to look up + // LOCAL_DISK replicas that belong to the client being drained. + std::unordered_map drain_segment_clients; + { + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + for (const auto& source_segment : job.request.segments) { + UUID client_id; + if (segment_access.GetClientIdBySegmentName(source_segment, + client_id) == + ErrorCode::OK) { + drain_segment_clients[source_segment] = client_id; + } + } + } + { std::shared_lock shared_lock(snapshot_mutex_); for (size_t i = 0; i < kNumShards; ++i) { MetadataShardAccessorRO shard(this, i); for (const auto& [tenant_id, tenant_state] : shard->tenants) { for (const auto& [key, metadata] : tenant_state.metadata) { - for (const auto& source_segment : job.request.segments) { - const auto unit_key = - MakeDrainUnitKey(tenant_id, key, source_segment); - if (job.completed_unit_keys.contains(unit_key) || - active_unit_keys.contains(unit_key) || - job.terminal_failed_unit_keys.contains(unit_key)) { - continue; - } - + for (const auto& source_segment : + job.request.segments) { + // Determine the move source: MEMORY segment name + // or LOCAL_DISK transport_endpoint. MEMORY takes + // priority (dedup: if both exist on the same + // client, only drain MEMORY). const auto replica_segments = metadata.GetReplicaSegmentNames(); - if (std::find(replica_segments.begin(), - replica_segments.end(), source_segment) == - replica_segments.end()) { + bool found_memory = + std::find(replica_segments.begin(), + replica_segments.end(), + source_segment) != + replica_segments.end(); + + std::string move_source = source_segment; + + if (!found_memory) { + // No MEMORY replica on this drain segment. + // Check if a LOCAL_DISK replica belongs to + // the client that owns this segment. + auto client_it = + drain_segment_clients.find(source_segment); + if (client_it == drain_segment_clients.end()) { + continue; + } + const auto& drain_client_id = + client_it->second; + + std::string transport_endpoint; + bool found_local_disk = false; + metadata.VisitReplicas( + [&drain_client_id](const Replica& r) { + return r.is_local_disk_replica() && + r.is_completed() && + r.get_local_disk_client_id() == + drain_client_id; + }, + [&transport_endpoint, &found_local_disk]( + const Replica& r) { + transport_endpoint = + r.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint; + found_local_disk = true; + }); + + if (!found_local_disk) { + continue; + } + + move_source = transport_endpoint; + } + + const auto unit_key = MakeDrainUnitKey( + tenant_id, key, move_source); + if (job.completed_unit_keys.contains(unit_key) || + active_unit_keys.contains(unit_key) || + job.terminal_failed_unit_keys.contains( + unit_key)) { continue; } if (metadata.IsHardPinned() || !metadata.IsLeaseExpired() || - !metadata.AllReplicas(&Replica::fn_is_completed) || + !metadata.AllReplicas( + &Replica::fn_is_completed) || tenant_state.replication_tasks.contains(key)) { blocked_unit_keys.insert(unit_key); continue; } + // Pass source_segment (not move_source) so that + // target selection avoids the source client's + // segments. auto target = SelectDrainTargetForKey( metadata, source_segment, job.request.target_segments); @@ -8595,8 +8677,9 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { } if (plans.size() < slots) { - plans.push_back({tenant_id, key, source_segment, - *target, metadata.size, unit_key}); + plans.push_back({tenant_id, key, move_source, + *target, metadata.size, + unit_key}); } } } @@ -8644,6 +8727,22 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { return false; } + // Pre-compute drain segment -> client_id mapping for LOCAL_DISK + // lookup. + std::unordered_map drain_segment_clients; + { + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + for (const auto& source_segment : job.request.segments) { + UUID client_id; + if (segment_access.GetClientIdBySegmentName(source_segment, + client_id) == + ErrorCode::OK) { + drain_segment_clients[source_segment] = client_id; + } + } + } + std::unordered_set remaining_segments; std::unordered_set remaining_unit_keys; { @@ -8655,6 +8754,7 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { const auto replica_segments = metadata.GetReplicaSegmentNames(); for (const auto& source_segment : job.request.segments) { + // MEMORY replica check (existing logic) if (std::find(replica_segments.begin(), replica_segments.end(), source_segment) != replica_segments.end()) { @@ -8662,6 +8762,35 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { remaining_unit_keys.insert(MakeDrainUnitKey( tenant_id, key, source_segment)); } + + // LOCAL_DISK replica check: if the client that + // owns this drain segment still has LOCAL_DISK + // data for this key, the segment is not fully + // drained yet. + auto client_it = + drain_segment_clients.find(source_segment); + if (client_it != drain_segment_clients.end()) { + const auto& drain_client_id = client_it->second; + metadata.VisitReplicas( + [&drain_client_id](const Replica& r) { + return r.is_local_disk_replica() && + r.is_completed() && + r.get_local_disk_client_id() == + drain_client_id; + }, + [&remaining_segments, &source_segment, + &tenant_id, &key, + &remaining_unit_keys, + this](const Replica& r) { + remaining_segments.insert(source_segment); + remaining_unit_keys.insert( + MakeDrainUnitKey( + tenant_id, key, + r.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint)); + }); + } } } } From 3f0bb081a09b83e4cd373d7878489fc00648fd9a Mon Sep 17 00:00:00 2001 From: liangxu Date: Thu, 23 Jul 2026 16:03:39 +0800 Subject: [PATCH 57/57] [Store] Complete LOCAL_DISK drain: gate LD metadata removal on drain moves, add e2e tests Completes the LOCAL_DISK drain migration (Option A) on top of the core implementation in 60e6d338: - MoveEnd drops redundant LOCAL_DISK metadata for drain moves only, gated by a new is_drain flag threaded through the move task payload (ReplicaMovePayload in task_manager.h), the client (Client::Move / ExecuteTask), the MasterClient::MoveEnd RPC and the rpc_service / master_service handler chain. Manual create_move_task defaults to is_drain=false so it no longer discards a live node's LOCAL_DISK backup. - Scope-narrowing fix for an EDEADLK on segment_mutex_ in MoveEnd (release the unique_lock before taking the shared_lock on the same shared_mutex). - Adds a replica_move_local_disk_transfer_success log in the LOCAL_DISK transfer branch of ExecuteReplicaTransfer, making the path observable and testable (file_storage.h: signature reformatting only). - mountLocalDiskSegment / getOffloadRpcAddr test helpers (real_client) and an offload_on_evict knob in the in-proc master test config. - Adds the local_disk_drain_test e2e suite (5 tests): baseline MEMORY-first dedup, is_drain gating regression (non-aligned sizes + position-dependent pattern), LOCAL_DISK load-failure graceful handling, and eviction-driven offload under both offload_on_evict modes. --- mooncake-store/include/client_service.h | 4 +- mooncake-store/include/file_storage.h | 5 +- mooncake-store/include/master_client.h | 3 +- mooncake-store/include/master_config.h | 8 + mooncake-store/include/master_service.h | 8 +- mooncake-store/include/real_client.h | 18 + mooncake-store/include/rpc_service.h | 3 +- mooncake-store/include/task_manager.h | 7 +- mooncake-store/src/client_service.cpp | 33 +- mooncake-store/src/master_client.cpp | 7 +- mooncake-store/src/master_service.cpp | 109 +- mooncake-store/src/real_client.cpp | 13 + mooncake-store/src/rpc_service.cpp | 11 +- mooncake-store/tests/e2e/CMakeLists.txt | 16 + .../tests/e2e/local_disk_drain_test.cpp | 1269 +++++++++++++++++ mooncake-store/tests/test_server_helpers.h | 3 + 16 files changed, 1451 insertions(+), 66 deletions(-) create mode 100644 mooncake-store/tests/e2e/local_disk_drain_test.cpp diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 156b9c8129..c8f51eea7b 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -32,7 +32,6 @@ namespace mooncake { class FileStorage; } - namespace mooncake { class PutOperation; @@ -914,7 +913,8 @@ class Client { tl::expected Move(const std::string& key, const std::string& tenant_id, const std::string& source, - const std::string& target); + const std::string& target, + bool is_drain = false); // Task thread pool for async task execution ThreadPool task_thread_pool_; diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index a0ccb71ddf..43906caf54 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -87,11 +87,10 @@ class FileStorage { * @return shared_ptr on success, error on failure */ tl::expected, ErrorCode> - LoadBatchFromLocalDisk(const std::string& key, - const std::string& tenant_id, uint64_t size); + LoadBatchFromLocalDisk(const std::string& key, const std::string& tenant_id, + uint64_t size); private: - /** * @brief Offload object data and metadata. * @return tl::expected indicating operation status. diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 7b54123044..900da6b827 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -549,7 +549,8 @@ class MasterClient { */ [[nodiscard]] tl::expected MoveEnd(const std::string& key); [[nodiscard]] tl::expected MoveEnd( - const std::string& key, const std::string& tenant_id); + const std::string& key, const std::string& tenant_id, + bool is_drain = false); /** * @brief Revoke a move operation diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index fd893953f8..db038574d7 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -1275,6 +1275,7 @@ struct InProcMasterConfig { std::optional root_fs_dir; std::optional enable_disk_eviction; std::optional quota_bytes; + std::optional offload_on_evict; }; // Builder class for InProcMasterConfig @@ -1292,6 +1293,7 @@ class InProcMasterConfigBuilder { std::optional root_fs_dir_ = std::nullopt; std::optional enable_disk_eviction_ = std::nullopt; std::optional quota_bytes_ = std::nullopt; + std::optional offload_on_evict_ = std::nullopt; public: InProcMasterConfigBuilder() = default; @@ -1360,6 +1362,11 @@ class InProcMasterConfigBuilder { return *this; } + InProcMasterConfigBuilder& set_offload_on_evict(bool enable) { + offload_on_evict_ = enable; + return *this; + } + InProcMasterConfig build() const; }; @@ -1378,6 +1385,7 @@ inline InProcMasterConfig InProcMasterConfigBuilder::build() const { config.root_fs_dir = root_fs_dir_; config.enable_disk_eviction = enable_disk_eviction_; config.quota_bytes = quota_bytes_; + config.offload_on_evict = offload_on_evict_; return config; } diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index e79f22ef84..cf18b82000 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -543,7 +543,8 @@ class MasterService { tl::expected MoveEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const std::string& tenant_id, + bool is_drain = false); tl::expected MoveRevoke(const UUID& client_id, const std::string& key, @@ -725,12 +726,15 @@ class MasterService { /** * @brief Create a move task to move an object's replica from source segment * to target segment + * @param is_drain True when the move is created by a drain job. Only drain + * moves drop the source client's redundant LOCAL_DISK metadata in MoveEnd. * @return Move task ID on success, ErrorCode on failure */ tl::expected CreateMoveTask(const std::string& key, const std::string& tenant_id, const std::string& source, - const std::string& target); + const std::string& target, + bool is_drain = false); /** * @brief Create a drain job to gracefully evacuate one or more segments. diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 870fa9ea7e..a545fd8050 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -723,6 +723,24 @@ class RealClient : public PyClient { int unmountSegment(const std::vector &segment_ids, uint64_t grace_period_seconds = 0); + /** + * @brief Mount a LOCAL_DISK (SSD offload) segment for this client on the + * master. Required before the master will allocate LOCAL_DISK + * replicas for this client's objects and push offload tasks to this + * client's FileStorage. Wraps Client::MountLocalDiskSegment. + * @param enable_offloading If true, enables offloading (write-to-file). + * @return 0 on success, negative value on error. + */ + int mountLocalDiskSegment(bool enable_offloading = true); + + /** + * @brief Offload RPC address used as the LOCAL_DISK replica segment name. + * This is the address the master records for LOCAL_DISK replicas of + * this client, and is the value to pass as the drain job's source + * segment when draining this client's LOCAL_DISK data. + */ + std::string getOffloadRpcAddr() const { return local_rpc_addr; } + /** * @brief Allocate memory internally and mount segments to master. * If size > max_mr_size, it will be split into multiple chunks. diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 8e2cfa4895..8da7a8116f 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -285,7 +285,8 @@ class WrappedMasterService { tl::expected MoveEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const std::string& tenant_id, + bool is_drain = false); tl::expected MoveRevoke(const UUID& client_id, const std::string& key, diff --git a/mooncake-store/include/task_manager.h b/mooncake-store/include/task_manager.h index 8f4578ce39..82edf9560b 100644 --- a/mooncake-store/include/task_manager.h +++ b/mooncake-store/include/task_manager.h @@ -104,8 +104,13 @@ struct ReplicaMovePayload { std::string key; std::string source; std::string target; + // True only when this move was created by a drain job. Drain moves drop + // the source client's redundant LOCAL_DISK metadata in MoveEnd; manual + // moves (e.g. the public create_move_task API) keep the LOCAL_DISK + // replica intact. + bool is_drain = false; }; -YLT_REFL(ReplicaMovePayload, tenant_id, key, source, target); +YLT_REFL(ReplicaMovePayload, tenant_id, key, source, target, is_drain); template struct TaskPayloadTraits; diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 55189ce617..da5fd562e2 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -3201,8 +3201,7 @@ tl::expected Client::ExecuteReplicaTransfer( file_storage_->LoadBatchFromLocalDisk(key, tenant_id, object_size); if (!load_result) { LOG(ERROR) << "action=replica_" << action_name << "_failed" - << ", key=" << key - << ", error=local_disk_load_failed" + << ", key=" << key << ", error=local_disk_load_failed" << ", error_code=" << load_result.error(); revoke_lambda(); return tl::unexpected(load_result.error()); @@ -3219,15 +3218,14 @@ tl::expected Client::ExecuteReplicaTransfer( auto slice_it = staging->slices.find(storage_key); if (slice_it == staging->slices.end()) { LOG(ERROR) << "action=replica_" << action_name << "_failed" - << ", key=" << key - << ", error=staging_slice_missing"; + << ", key=" << key << ", error=staging_slice_missing"; revoke_lambda(); return tl::unexpected(ErrorCode::INTERNAL_ERROR); } // (b) Split the staging buffer into transfer slices. - auto slices = split_into_slices(slice_it->second.ptr, - slice_it->second.size); + auto slices = + split_into_slices(slice_it->second.ptr, slice_it->second.size); // (c) TransferWrite to each target. for (const auto& target : targets) { @@ -3240,6 +3238,14 @@ tl::expected Client::ExecuteReplicaTransfer( } } + // LOCAL_DISK transfer succeeded: data flowed local-disk SSD -> staging + // buffer (LoadBatchFromLocalDisk) -> TransferWrite -> target MEMORY. + // Logged so the LOCAL_DISK drain path is observable and testable. + LOG(INFO) << "action=replica_" << action_name + << "_local_disk_transfer_success" + << ", key=" << key << ", object_size=" << object_size + << ", path=local_disk_ssd->staging_buffer->transfer_write"; + // (d) Finalize. The staging buffer is released when staging goes // out of scope at the end of this block. auto end_result = end_fn(); @@ -3354,9 +3360,11 @@ tl::expected Client::Move(const std::string& key, tl::expected Client::Move(const std::string& key, const std::string& tenant_id, const std::string& source, - const std::string& target) { + const std::string& target, + bool is_drain) { LOG(INFO) << "action=replica_move_start" << ", key=" << key - << ", source_segment=" << source << ", target_segment=" << target; + << ", source_segment=" << source << ", target_segment=" << target + << ", is_drain=" << is_drain; // Call MoveStart first - it validates existence and allocates replica if // needed @@ -3375,7 +3383,7 @@ tl::expected Client::Move(const std::string& key, LOG(INFO) << "action=replica_move_skipped" << ", key=" << key << ", info=target_replica_already_exists"; // Target already exists, consider it success - auto move_end_result = master_client_.MoveEnd(key, tenant_id); + auto move_end_result = master_client_.MoveEnd(key, tenant_id, is_drain); if (!move_end_result.has_value()) { ErrorCode error = move_end_result.error(); LOG(ERROR) << "action=replica_move_failed" << ", key=" << key @@ -3389,7 +3397,7 @@ tl::expected Client::Move(const std::string& key, auto result = ExecuteReplicaTransfer( key, tenant_id, "move", - [&]() { return master_client_.MoveEnd(key, tenant_id); }, + [&]() { return master_client_.MoveEnd(key, tenant_id, is_drain); }, [&]() { return master_client_.MoveRevoke(key, tenant_id); }, response.source, targets); @@ -3681,8 +3689,9 @@ void Client::ExecuteTask(const ClientTask& client_task) { case TaskType::REPLICA_MOVE: { ReplicaMovePayload payload; struct_json::from_json(payload, assignment.payload); - auto move_result = Move(payload.key, payload.tenant_id, - payload.source, payload.target); + auto move_result = + Move(payload.key, payload.tenant_id, payload.source, + payload.target, payload.is_drain); if (move_result.has_value()) { result = ErrorCode::OK; } else { diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index e2d9db347f..0997dc086e 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -1190,12 +1190,13 @@ tl::expected MasterClient::MoveEnd(const std::string& key) { } tl::expected MasterClient::MoveEnd( - const std::string& key, const std::string& tenant_id) { + const std::string& key, const std::string& tenant_id, bool is_drain) { ScopedVLogTimer timer(1, "MasterClient::MoveEnd"); - timer.LogRequest("key=", key, ", tenant_id=", tenant_id); + timer.LogRequest("key=", key, ", tenant_id=", tenant_id, + ", is_drain=", is_drain); auto result = invoke_rpc<&WrappedMasterService::MoveEnd, void>( - client_id_, key, tenant_id); + client_id_, key, tenant_id, is_drain); timer.LogResponseExpected(result); return result; } diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index d1d1ec3944..ff4c8c9267 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -4336,8 +4336,8 @@ tl::expected MasterService::MoveStart( } tl::expected MasterService::MoveEnd( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const std::string& tenant_id, + bool is_drain) { std::shared_lock shared_lock(snapshot_mutex_); MetadataAccessorRW accessor(this, MakeObjectIdentityForRequest(key, tenant_id)); @@ -4423,6 +4423,50 @@ tl::expected MasterService::MoveEnd( return replica.id() == source_id; }); if (!source_replica.empty()) { + // Drain-only cleanup: when a *drain* move relocates a MEMORY replica, + // any LOCAL_DISK replica of the same key owned by the same client is + // now redundant (it is a disk backup of the memory data we just + // relocated, and the target already holds that data). Drop the LD + // replica's master metadata so the source client is fully drained. + // The underlying SSD file is intentionally left on disk as a safety + // net (no physical deletion); physical cleanup is a separate, future + // step. See drain-local-disk design discussion. + // + // This is gated on is_drain so that a *manual* move (e.g. the public + // create_move_task API used for load-balancing) does NOT silently + // discard the source client's LOCAL_DISK backup or leak its SSD file + // on a still-live node. + if (is_drain && source_replica.front().is_memory_replica()) { + const auto& mem_segments = + source_replica.front().get_segment_names(); + if (!mem_segments.empty() && mem_segments.front().has_value()) { + UUID source_client_id; + // Narrow scope: release segment_mutex_ (unique_lock) before + // EraseReplicasWithCacheTotalAccounting, which internally + // takes segment_mutex_ shared_lock via + // getLocalDiskSegmentAccess(). Holding unique then taking + // shared on the same std::shared_mutex in the same thread + // triggers libstdc++ shared_mutex::lock_shared to throw + // system_error(EDEADLK). + bool found_source_client = false; + { + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + found_source_client = + (segment_access.GetClientIdBySegmentName( + mem_segments.front().value(), source_client_id) == + ErrorCode::OK); + } + if (found_source_client) { + EraseReplicasWithCacheTotalAccounting( + metadata, [&source_client_id](const Replica& replica) { + return replica.is_local_disk_replica() && + replica.get_local_disk_client_id() == + source_client_id; + }); + } + } + } std::lock_guard lock(discarded_replicas_mutex_); discarded_replicas_.emplace_back( std::move(source_replica), @@ -8176,7 +8220,7 @@ tl::expected MasterService::CreateCopyTask( tl::expected MasterService::CreateMoveTask( const std::string& key, const std::string& tenant_id, - const std::string& source, const std::string& target) { + const std::string& source, const std::string& target, bool is_drain) { auto normalized_tenant_result = NormalizeTenantIdForWrite(tenant_id); if (!normalized_tenant_result) { return tl::make_unexpected(normalized_tenant_result.error()); @@ -8228,8 +8272,7 @@ tl::expected MasterService::CreateMoveTask( if (error != ErrorCode::OK) { auto local_disk_replica = metadata.GetReplicaBySegmentName(source); if (local_disk_replica && local_disk_replica->is_local_disk_replica()) { - auto client_id_opt = - local_disk_replica->get_local_disk_client_id(); + auto client_id_opt = local_disk_replica->get_local_disk_client_id(); if (client_id_opt.has_value()) { select_client = client_id_opt.value(); error = ErrorCode::OK; @@ -8248,7 +8291,8 @@ tl::expected MasterService::CreateMoveTask( select_client, {.tenant_id = object_id.tenant_id, .key = object_id.user_key, .source = source, - .target = target}); + .target = target, + .is_drain = is_drain}); } tl::expected MasterService::QueryTask( @@ -8580,9 +8624,8 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { segment_manager_.getSegmentAccess(); for (const auto& source_segment : job.request.segments) { UUID client_id; - if (segment_access.GetClientIdBySegmentName(source_segment, - client_id) == - ErrorCode::OK) { + if (segment_access.GetClientIdBySegmentName( + source_segment, client_id) == ErrorCode::OK) { drain_segment_clients[source_segment] = client_id; } } @@ -8594,8 +8637,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { MetadataShardAccessorRO shard(this, i); for (const auto& [tenant_id, tenant_state] : shard->tenants) { for (const auto& [key, metadata] : tenant_state.metadata) { - for (const auto& source_segment : - job.request.segments) { + for (const auto& source_segment : job.request.segments) { // Determine the move source: MEMORY segment name // or LOCAL_DISK transport_endpoint. MEMORY takes // priority (dedup: if both exist on the same @@ -8605,8 +8647,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { bool found_memory = std::find(replica_segments.begin(), replica_segments.end(), - source_segment) != - replica_segments.end(); + source_segment) != replica_segments.end(); std::string move_source = source_segment; @@ -8619,8 +8660,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { if (client_it == drain_segment_clients.end()) { continue; } - const auto& drain_client_id = - client_it->second; + const auto& drain_client_id = client_it->second; std::string transport_endpoint; bool found_local_disk = false; @@ -8631,8 +8671,8 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { r.get_local_disk_client_id() == drain_client_id; }, - [&transport_endpoint, &found_local_disk]( - const Replica& r) { + [&transport_endpoint, + &found_local_disk](const Replica& r) { transport_endpoint = r.get_descriptor() .get_local_disk_descriptor() @@ -8647,19 +8687,17 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { move_source = transport_endpoint; } - const auto unit_key = MakeDrainUnitKey( - tenant_id, key, move_source); + const auto unit_key = + MakeDrainUnitKey(tenant_id, key, move_source); if (job.completed_unit_keys.contains(unit_key) || active_unit_keys.contains(unit_key) || - job.terminal_failed_unit_keys.contains( - unit_key)) { + job.terminal_failed_unit_keys.contains(unit_key)) { continue; } if (metadata.IsHardPinned() || !metadata.IsLeaseExpired() || - !metadata.AllReplicas( - &Replica::fn_is_completed) || + !metadata.AllReplicas(&Replica::fn_is_completed) || tenant_state.replication_tasks.contains(key)) { blocked_unit_keys.insert(unit_key); continue; @@ -8678,8 +8716,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { if (plans.size() < slots) { plans.push_back({tenant_id, key, move_source, - *target, metadata.size, - unit_key}); + *target, metadata.size, unit_key}); } } } @@ -8691,7 +8728,8 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { for (const auto& plan : plans) { auto task_id = CreateMoveTask(plan.key, plan.tenant_id, - plan.source_segment, plan.target_segment); + plan.source_segment, plan.target_segment, + /*is_drain=*/true); if (task_id.has_value()) { ActiveDrainTask active_task; active_task.task_id = task_id.value(); @@ -8735,9 +8773,8 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { segment_manager_.getSegmentAccess(); for (const auto& source_segment : job.request.segments) { UUID client_id; - if (segment_access.GetClientIdBySegmentName(source_segment, - client_id) == - ErrorCode::OK) { + if (segment_access.GetClientIdBySegmentName( + source_segment, client_id) == ErrorCode::OK) { drain_segment_clients[source_segment] = client_id; } } @@ -8779,16 +8816,14 @@ bool MasterService::MaybeCompleteDrainJob(DrainJob& job) { drain_client_id; }, [&remaining_segments, &source_segment, - &tenant_id, &key, - &remaining_unit_keys, + &tenant_id, &key, &remaining_unit_keys, this](const Replica& r) { remaining_segments.insert(source_segment); - remaining_unit_keys.insert( - MakeDrainUnitKey( - tenant_id, key, - r.get_descriptor() - .get_local_disk_descriptor() - .transport_endpoint)); + remaining_unit_keys.insert(MakeDrainUnitKey( + tenant_id, key, + r.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint)); }); } } diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index fe00d17224..f9aa966bdc 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1410,6 +1410,19 @@ int RealClient::unmountSegment(const std::vector &segment_ids, return first_error; } +int RealClient::mountLocalDiskSegment(bool enable_offloading) { + if (!client_) { + LOG(ERROR) << "Client not initialized"; + return -1; + } + auto result = client_->MountLocalDiskSegment(enable_offloading); + if (!result) { + LOG(ERROR) << "MountLocalDiskSegment failed: " << result.error(); + return -1; + } + return 0; +} + int RealClient::allocateAndMountSegment( size_t size, const std::string &protocol, const std::string &location, std::vector &out_segment_ids, size_t *out_allocated_size) { diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index d5de362f54..820c58fc1a 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -891,14 +891,17 @@ tl::expected WrappedMasterService::MoveStart( } tl::expected WrappedMasterService::MoveEnd( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const std::string& tenant_id, + bool is_drain) { return execute_rpc( "MoveEnd", - [&] { return master_service_.MoveEnd(client_id, key, tenant_id); }, + [&] { + return master_service_.MoveEnd(client_id, key, tenant_id, is_drain); + }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, - ", tenant_id=", tenant_id); + ", tenant_id=", tenant_id, + ", is_drain=", is_drain); }, [] { MasterMetricManager::instance().inc_move_end_requests(); }, [] { MasterMetricManager::instance().inc_move_end_failures(); }); diff --git a/mooncake-store/tests/e2e/CMakeLists.txt b/mooncake-store/tests/e2e/CMakeLists.txt index ce23d3062b..87e19297a5 100644 --- a/mooncake-store/tests/e2e/CMakeLists.txt +++ b/mooncake-store/tests/e2e/CMakeLists.txt @@ -81,3 +81,19 @@ target_link_libraries(storage_backend_e2e_test PUBLIC ) add_test(NAME storage_backend_e2e_test COMMAND storage_backend_e2e_test) + +add_executable(local_disk_drain_test local_disk_drain_test.cpp) +target_include_directories(local_disk_drain_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/.. +) +target_link_libraries(local_disk_drain_test PUBLIC + mooncake_store + transfer_engine + cachelib_memory_allocator + glog + gtest + pthread + ${ETCD_WRAPPER_LIB} +) + +add_test(NAME local_disk_drain_test COMMAND local_disk_drain_test) diff --git a/mooncake-store/tests/e2e/local_disk_drain_test.cpp b/mooncake-store/tests/e2e/local_disk_drain_test.cpp new file mode 100644 index 0000000000..53d3898621 --- /dev/null +++ b/mooncake-store/tests/e2e/local_disk_drain_test.cpp @@ -0,0 +1,1269 @@ +// Copyright 2025 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. + +// ------------------------------------------------------------------- +// LocalDiskDrainTest +// +// In-process master (non-HA) + two offload-capable RealClients over +// TCP that verifies the LOCAL_DISK drain job path introduced by the +// "LOCAL_DISK drain migration via Move flow (Option A)" change: +// +// Put (MEMORY + LOCAL_DISK replica on client1's SSD) +// -> drain job moves the LOCAL_DISK replica off client1's segment +// -> data is now served from client2, byte-for-byte identical +// +// This complements task_integration_test.DrainJobCompleteFlow, which +// only exercises in-memory (DRAM) segment drains. Here the source +// replica is physically written to the client's local disk +// (FileStorage) and the drain job must read it back via the new +// LOCAL_DISK branch of Client::ExecuteReplicaTransfer. +// ------------------------------------------------------------------- + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "real_client.h" +#include "rpc_types.h" +#include "master_metric_manager.h" +#include "test_server_helpers.h" +#include "types.h" +#include "utils.h" + +DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); +DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma"); + +namespace fs = std::filesystem; + +namespace mooncake { +namespace testing { + +// ----- HTTP drain-job API response structs (mirrors master_admin_service) +// ----- +struct HttpCreateDrainJobResponse { + bool success{false}; + std::string job_id; + std::string status; + int32_t error_code{0}; + std::string error_message; +}; +YLT_REFL(HttpCreateDrainJobResponse, success, job_id, status, error_code, + error_message); + +struct HttpQueryDrainJobResponse { + bool success{false}; + std::string job_id; + int32_t type{0}; + std::string type_name; + int32_t status{0}; + std::string status_name; + int64_t created_at_ms_epoch{0}; + int64_t last_updated_at_ms_epoch{0}; + std::vector segments; + uint64_t succeeded_units{0}; + uint64_t failed_units{0}; + uint64_t blocked_units{0}; + uint64_t active_units{0}; + uint64_t migrated_bytes{0}; + std::string message; + int32_t error_code{0}; + std::string error_message; +}; +YLT_REFL(HttpQueryDrainJobResponse, success, job_id, type, type_name, status, + status_name, created_at_ms_epoch, last_updated_at_ms_epoch, segments, + succeeded_units, failed_units, blocked_units, active_units, + migrated_bytes, message, error_code, error_message); + +struct HttpSegmentStatusResponse { + bool success{false}; + std::string segment; + int32_t status{0}; + std::string status_name; + int32_t error_code{0}; + std::string error_message; +}; +YLT_REFL(HttpSegmentStatusResponse, success, segment, status, status_name, + error_code, error_message); + +tl::expected HttpPostJson(const std::string& url, + const std::string& body) { + coro_http::coro_http_client client; + auto response = client.post(url, body, coro_http::req_content_type::json); + if (response.status != 200) { + return tl::unexpected(response.status); + } + return std::string(response.resp_body); +} + +tl::expected HttpGet(const std::string& url) { + coro_http::coro_http_client client; + auto response = client.get(url); + if (response.status != 200) { + return tl::unexpected(response.status); + } + return std::string(response.resp_body); +} + +// Build a deterministic value whose bytes vary with BOTH position and a +// per-key seed. A byte-for-byte comparison of the drained data then detects +// intra-object slice mis-ordering or a wrong O_DIRECT read offset -- bugs a +// single-repeated-byte fill (e.g. string(size, 'A')) would silently pass. +std::string MakePatternValue(size_t size, uint32_t seed) { + std::string v; + v.resize(size); + for (size_t j = 0; j < size; ++j) { + v[j] = static_cast((j * 31u + seed * 131u + 7u) & 0xFFu); + } + return v; +} + +class LocalDiskDrainTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("LocalDiskDrainTest"); + FLAGS_logtostderr = 1; + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + void SetUp() override { + tmp_root_ = fs::temp_directory_path() / + ("mc_lddrain_" + std::to_string(::getpid())); + fs::create_directories(tmp_root_); + + // Force the client-side BucketStorageBackend to seal (flush) a bucket + // to disk after every single offloaded object. By default a bucket + // only flushes once it accumulates 500 keys or 256 MB, so the few + // small objects this test writes would stay buffered in memory and + // never appear on disk, making the LOCAL_DISK file assertion time out. + // Setting the key limit to 1 makes each offloaded object flush + // immediately, so the offload actually materializes on the client SSD. + ::setenv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "1", /*overwrite=*/1); + + // Master must have offload enabled so the clients can offload their + // MEMORY replicas to LOCAL_DISK (client-side SSD). We intentionally do + // NOT set root_fs_dir: that would enable the master's *server-side* + // DISK (NoF) offload, which writes DISK replicas under + // root_fs_dir/mooncake_cluster and would shadow the LOCAL_DISK path + // this test exercises. + InProcMasterConfig config = + InProcMasterConfigBuilder().set_enable_offload(true).build(); + ASSERT_TRUE(master_.Start(config)) + << "Failed to start in-proc master with offload enabled"; + master_address_ = master_.master_address(); + LOG(INFO) << "Started in-proc master at " << master_address_; + } + + void TearDown() override { + if (client1_) client1_->tearDownAll(); + if (client2_) client2_->tearDownAll(); + if (client3_) client3_->tearDownAll(); + master_.Stop(); + std::error_code ec; + fs::remove_all(tmp_root_, ec); + } + + void SetupOffloadClient(const std::string& addr, const fs::path& ssd_path, + std::shared_ptr& out) { + fs::create_directories(ssd_path); + const std::string rdma_devices = + (FLAGS_protocol == "rdma") ? FLAGS_device_name : std::string(""); + out = RealClient::create(); + ASSERT_NE(out, nullptr); + constexpr size_t kGlobalSegmentSize = 64 * 1024 * 1024; + constexpr size_t kLocalBufferSize = 64 * 1024 * 1024; + ASSERT_EQ( + out->setup_real(addr, "P2PHANDSHAKE", + /*global_segment_size=*/kGlobalSegmentSize, + /*local_buffer_size=*/kLocalBufferSize, + FLAGS_protocol, rdma_devices, master_address_, + /*transfer_engine=*/nullptr, + /*ipc_socket_path=*/"", + /*enable_ssd_offload=*/true, ssd_path.string(), + /*tenant_id=*/"default"), + 0) + << "Failed to setup offload client at " << addr; + // Config trace: MEMORY segment + LOCAL_DISK SSD path/capacity. + // total_size_limit for the SSD backend defaults to 2 TB + // (storage_backend.h kDefault) when not explicitly overridden, + // which is effectively unbounded for this test's 16 KB payload. + LOG(INFO) << "[Setup] client=" << addr + << " MEMORY_segment_size=" << kGlobalSegmentSize << " (" + << kGlobalSegmentSize / 1024 << " KiB)" + << " local_buffer_size=" << kLocalBufferSize << " (" + << kLocalBufferSize / 1024 << " KiB)" + << " SSD_offload_path=" << ssd_path.string() + << " SSD_total_size_limit=default(2TB)"; + } + + tl::expected CreateDrainJobViaHttp( + const CreateDrainJobRequest& request) { + std::string body; + struct_json::to_json(request, body); + auto response = HttpPostJson( + master_.http_metrics_base() + "/api/v1/drain_jobs", body); + if (!response.has_value()) return tl::unexpected(response.error()); + HttpCreateDrainJobResponse parsed; + struct_json::from_json(parsed, response.value()); + return parsed; + } + + tl::expected QueryDrainJobViaHttp( + const std::string& job_id) { + auto response = HttpGet(master_.http_metrics_base() + + "/api/v1/drain_jobs/query?job_id=" + job_id); + if (!response.has_value()) return tl::unexpected(response.error()); + HttpQueryDrainJobResponse parsed; + struct_json::from_json(parsed, response.value()); + return parsed; + } + + tl::expected QuerySegmentStatusViaHttp( + const std::string& segment_name) { + auto response = + HttpGet(master_.http_metrics_base() + + "/api/v1/segments/status?segment=" + segment_name); + if (!response.has_value()) return tl::unexpected(response.error()); + HttpSegmentStatusResponse parsed; + struct_json::from_json(parsed, response.value()); + return parsed; + } + + bool WaitForJobCompletionViaHttp( + const std::string& job_id, HttpQueryDrainJobResponse* final_job, + std::chrono::seconds timeout = std::chrono::seconds(120)) { + auto start = std::chrono::steady_clock::now(); + while (std::chrono::steady_clock::now() - start < timeout) { + auto query_result = QueryDrainJobViaHttp(job_id); + if (query_result.has_value()) { + if (query_result->status_name == "SUCCEEDED" || + query_result->status_name == "FAILED" || + query_result->status_name == "CANCELED") { + if (final_job != nullptr) *final_job = query_result.value(); + return query_result->status_name == "SUCCEEDED"; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + return false; + } + + // Poll until the local-disk directory contains at least one regular file. + // This proves the offload path actually wrote LOCAL_DISK replica data to + // the client's SSD. We cannot use get_replica_desc()/Query() for this: + // that RPC is for transfer planning and intentionally returns only + // transferable replicas (MEMORY + DISK/NoF), never LOCAL_DISK replicas. + bool WaitForLocalDiskFiles( + const fs::path& ssd_path, + std::chrono::seconds timeout = std::chrono::seconds(120)) { + auto deadline = std::chrono::steady_clock::now() + timeout; + int poll = 0; + while (std::chrono::steady_clock::now() < deadline) { + std::error_code ec; + bool has_files = false; + for (auto it = fs::recursive_directory_iterator(ssd_path, ec); + it != fs::recursive_directory_iterator(); ++it) { + if (it->is_regular_file()) { + has_files = true; + break; + } + } + if (has_files) return true; + if ((++poll % 10) == 0) { + LOG(INFO) << "[WaitForLocalDiskFiles] poll=" << poll + << " ssd_path=" << ssd_path.string() + << " still empty"; + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + LOG(ERROR) << "[WaitForLocalDiskFiles] TIMEOUT: no LOCAL_DISK files in " + << ssd_path.string(); + // Dump the whole temp tree so we can see where the offloaded files + // actually landed. + { + std::error_code ec; + std::string tree; + for (auto it = fs::recursive_directory_iterator(tmp_root_, ec); + it != fs::recursive_directory_iterator(); ++it) { + tree += "\n " + it->path().string(); + } + LOG(ERROR) << "[WaitForLocalDiskFiles] tmp_root tree:" << tree; + } + return false; + } + + // Count regular files and total bytes under an SSD offload directory. + // Used to leave a trace of how many LOCAL_DISK replica files physically + // landed on each client's SSD before/after drain. + struct SsdFileStats { + size_t file_count{0}; + uint64_t total_bytes{0}; + }; + SsdFileStats CountSsdFiles(const fs::path& ssd_path) { + SsdFileStats stats; + std::error_code ec; + for (auto it = fs::recursive_directory_iterator(ssd_path, ec); + it != fs::recursive_directory_iterator(); ++it) { + if (it->is_regular_file()) { + stats.file_count++; + stats.total_bytes += it->file_size(ec); + } + } + return stats; + } + + InProcMaster master_; + std::string master_address_; + std::shared_ptr client1_; + std::shared_ptr client2_; + std::shared_ptr client3_; + fs::path tmp_root_; +}; + +TEST_F(LocalDiskDrainTest, DrainLocalDiskReplicasToPeer) { + const std::string client1_addr = "127.0.0.1:17813"; + const std::string client2_addr = "127.0.0.1:17814"; + fs::path ssd1 = tmp_root_ / "client1_ssd"; + fs::path ssd2 = tmp_root_ / "client2_ssd"; + + SetupOffloadClient(client1_addr, ssd1, client1_); + SetupOffloadClient(client2_addr, ssd2, client2_); + + // Register a LOCAL_DISK (SSD offload) segment for client1 on the master. + // Without this, the master never allocates LOCAL_DISK replicas for + // client1's objects and never pushes offload tasks to client1's + // FileStorage, so nothing is ever written to ssd1. This is the missing + // piece that makes this test exercise real LOCAL_DISK data, not just + // MEMORY replicas. + ASSERT_EQ(client1_->mountLocalDiskSegment(true), 0) + << "Failed to mount LOCAL_DISK segment on client1"; + + // Phase 1: Put keys on client1. With offload enabled the master assigns + // LOCAL_DISK replicas and client1 writes them to its local disk. + // CRITICAL: set preferred_segment = client1_addr so the master allocator + // places the MEMORY replica on client1 (not client2). Without this, the + // allocator may spread replicas across both clients, causing only a subset + // of keys to be on the drain source and making succeeded_units flaky. + constexpr int kKeyCount = 4; + constexpr size_t kValueSize = 4096; + constexpr const char* kSourceSegment = "127.0.0.1:17813"; + std::vector keys; + std::vector values; + for (int i = 0; i < kKeyCount; ++i) { + std::string key = "ld_drain_key_" + std::to_string(i); + std::string value = MakePatternValue(kValueSize, i); + std::span span(value.data(), value.size()); + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = kSourceSegment; + ASSERT_EQ(client1_->put(key, span, config), 0) + << "Put failed for " << key; + keys.push_back(key); + values.push_back(value); + } + // Data trace: total payload and per-replica footprint. + // With offload enabled, each key produces 1 MEMORY + 1 LOCAL_DISK + // replica on client1, so the on-wire footprint is 2x payload. + LOG(INFO) << "[Phase 1] Put " << kKeyCount << " keys x " << kValueSize + << " bytes = " << (uint64_t)kKeyCount * kValueSize << " bytes (" + << (uint64_t)kKeyCount * kValueSize / 1024 << " KiB)" + << " preferred_segment=" << kSourceSegment + << "; with offload, client1 holds MEMORY+LOCAL_DISK replicas = " + << 2 * (uint64_t)kKeyCount * kValueSize << " bytes (" + << 2 * (uint64_t)kKeyCount * kValueSize / 1024 << " KiB)"; + + // Phase 2: Wait for the LOCAL_DISK replicas to be physically written to + // client1's SSD. This is the core assertion that distinguishes this test + // from the in-memory drain test. We verify it via the filesystem (the + // offload path writes object files under ssd1), because get_replica_desc/ + // Query intentionally does NOT expose LOCAL_DISK replicas. + ASSERT_TRUE(WaitForLocalDiskFiles(ssd1, std::chrono::seconds(120))) + << "LOCAL_DISK replicas were not offloaded to client1's disk"; + + // Trace: confirm how many LOCAL_DISK files landed on client1's SSD + // before drain, and that client2's SSD is still empty. + { + const auto stats1 = CountSsdFiles(ssd1); + const auto stats2 = CountSsdFiles(ssd2); + LOG(INFO) << "[Phase 2] Pre-drain SSD footprint:" + << " client1_ssd files=" << stats1.file_count + << " bytes=" << stats1.total_bytes + << " | client2_ssd files=" << stats2.file_count + << " bytes=" << stats2.total_bytes; + } + + // Wait for the KV lease to expire before creating the drain job. + // ScheduleDrainJobTasks blocks any key whose lease is still active + // (master_service.cpp ScheduleDrainJobTasks: IsHardPinned || + // !IsLeaseExpired || !AllReplicas || has_repl_task). The default TTL is + // DEFAULT_DEFAULT_KV_LEASE_TTL = 5000ms (types.h). Sleeping 6s guarantees + // all 4 keys are eligible on the first scheduling pass, so + // succeeded_units converges to kKeyCount instead of fluctuating. + std::this_thread::sleep_for(std::chrono::seconds(6)); + + // The drain job is keyed by the client's *MEMORY* segment name (the + // address the segment is registered under on the master), NOT the offload + // RPC address. CreateDrainJob validates the source against the registered + // segment table and sets it to DRAINING; the offload RPC address is not a + // registered segment and would be rejected. The drain scheduler maps this + // MEMORY segment to its owning client_id and, once the MEMORY replica has + // moved, drains that client's LOCAL_DISK replica by its transport_endpoint + // internally (see MasterService drain Option A logic). + const std::string source_segment = client1_addr; + const std::string target_segment = client2_addr; + LOG(INFO) << "Draining source_segment=" << source_segment + << " -> target_segment=" << target_segment; + + // Phase 3: Sanity check - data is readable on the source before drain. + { + auto alloc = client1_->client_buffer_allocator_->allocate(kValueSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client1_->get_into(keys[0], dst, kValueSize); + ASSERT_EQ(n, static_cast(kValueSize)) + << "Pre-drain get from source failed"; + ASSERT_EQ(std::string(dst, kValueSize), values[0]); + LOG(INFO) << "[Phase 3] Sanity check OK: key=" << keys[0] + << " readable from client1 (pre-drain)"; + } + + // Phase 4: Create the drain job over the master admin HTTP API. + CreateDrainJobRequest request; + request.segments = {source_segment}; + request.target_segments = {target_segment}; + request.max_concurrency = 1; + auto create_res = CreateDrainJobViaHttp(request); + ASSERT_TRUE(create_res.has_value()) << "CreateDrainJob HTTP call failed"; + ASSERT_TRUE(create_res->success) << create_res->error_message; + EXPECT_EQ(create_res->status, "CREATED"); + + // Phase 5: Wait for the drain job to finish. + HttpQueryDrainJobResponse final_job; + ASSERT_TRUE(WaitForJobCompletionViaHttp(create_res->job_id, &final_job, + std::chrono::seconds(120))) + << "Drain job did not complete (status=" + << (final_job.status_name.empty() ? "" : final_job.status_name) + << ")"; + EXPECT_EQ(final_job.status_name, "SUCCEEDED"); + EXPECT_EQ(final_job.failed_units, 0u); + // All kKeyCount keys must have been migrated. With preferred_segment set + // and lease expired, ScheduleDrainJobTasks should not block or skip any + // key. A lower value would indicate a regression (e.g. lease not expired, + // allocator scattered replicas, or MaybeCompleteDrainJob terminating + // early). + EXPECT_EQ(final_job.succeeded_units, static_cast(kKeyCount)) + << "Expected all " << kKeyCount + << " keys to be migrated, got succeeded_units=" + << final_job.succeeded_units; + // Trace: drain result summary from master's POV. + LOG(INFO) << "[Phase 5] Drain job " << final_job.job_id + << " status=" << final_job.status_name + << " succeeded_units=" << final_job.succeeded_units + << " failed_units=" << final_job.failed_units + << " blocked_units=" << final_job.blocked_units + << " migrated_bytes=" << final_job.migrated_bytes + << " message=" << final_job.message; + + // The source segment must transition to DRAINED, which only happens when + // the master sees *no* replicas (neither MEMORY nor LOCAL_DISK) belonging + // to the drained client. This proves the redundant LOCAL_DISK metadata was + // dropped on the source (its SSD file is intentionally retained as a safety + // net) and the client is fully drained. + { + auto drained_status = QuerySegmentStatusViaHttp(source_segment); + ASSERT_TRUE(drained_status.has_value()) + << "Failed to query source segment status over HTTP"; + EXPECT_TRUE(drained_status->success); + EXPECT_EQ(drained_status->status_name, "DRAINED"); + } + + // Phase 6: After drain, every key must be byte-for-byte readable from the + // target client (data survived the LOCAL_DISK -> peer migration). + for (int i = 0; i < kKeyCount; ++i) { + auto alloc = client2_->client_buffer_allocator_->allocate(kValueSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client2_->get_into(keys[i], dst, kValueSize); + ASSERT_EQ(n, static_cast(kValueSize)) + << "Get from target failed for " << keys[i]; + ASSERT_EQ(std::string(dst, kValueSize), values[i]) + << "Data mismatch after drain for " << keys[i]; + } + // Trace: confirm migration destination = client2 MEMORY. + // Data moved: 4 keys x 4 KiB = 16 KiB from client1 (MEMORY+LOCAL_DISK) + // to client2 (MEMORY only, per MEMORY-first dedup design). + { + const auto stats1 = CountSsdFiles(ssd1); + const auto stats2 = CountSsdFiles(ssd2); + LOG(INFO) << "[Phase 6] Post-drain footprint:" + << " client1_ssd residual_files=" << stats1.file_count + << " bytes=" << stats1.total_bytes + << " (safety-net retained, expected non-zero)" + << " | client2_ssd files=" << stats2.file_count + << " bytes=" << stats2.total_bytes + << " (expected 0: target stores MEMORY replicas only)"; + LOG(INFO) + << "[Phase 6] Data migrated: " << kKeyCount << " keys x " + << kValueSize << " bytes = " << (uint64_t)kKeyCount * kValueSize + << " bytes (" << (uint64_t)kKeyCount * kValueSize / 1024 << " KiB)" + << " from client1 MEMORY -> client2 MEMORY" + << " (LOCAL_DISK dedup: only MEMORY moved, LD metadata dropped)"; + } + + // Phase 7: Source-side SSD cleanup after LOCAL_DISK migration is a known + // follow-up (see RFC: "Source SSD cleanup ... residual SSD data is + // harmless"). So we only log whether residual files remain; we do NOT + // assert their absence. The correctness guarantee verified above is that + // the LOCAL_DISK replica was migrated (drain SUCCEEDED) and the data is + // byte-for-byte readable from the target client. + { + const auto stats1 = CountSsdFiles(ssd1); + LOG(INFO) + << "[Phase 7] Source client1 residual LOCAL_DISK files on ssd1: " + << stats1.file_count << " files / " << stats1.total_bytes + << " bytes " + << (stats1.file_count > 0 + ? "(present, expected - safety net, not cleaned)" + : "(none)"); + } +} + +// ------------------------------------------------------------------- +// LocalDiskDrainEvictTest +// +// Large-data eviction tests: fill a 64 MB MEMORY segment with 4 MB +// objects, trigger eviction via allocation failure, then drain the +// resulting LOCAL_DISK replicas to a peer client. +// +// Two scenarios: +// 1. offload_on_evict=true 鈥?offload deferred to eviction time +// 2. offload_on_evict=false 鈥?offload immediate at PutEnd +// ------------------------------------------------------------------- +// A glog sink that captures log messages in-process so the test can assert +// that the LOCAL_DISK drain path emitted its success log +// (action=replica_move_local_disk_transfer_success). This is DIRECT evidence +// that data flowed local-disk SSD -> staging buffer -> TransferWrite, rather +// than an inference from the memory ratio. +class CapturingLogSink : public google::LogSink { + public: + void send(google::LogSeverity /*severity*/, const char* /*full_filename*/, + const char* /*base_filename*/, int /*line*/, + const struct ::tm* /*tm_time*/, const char* message, + size_t message_len) override { + std::lock_guard lock(mutex_); + messages_.emplace_back(message, message_len); + } + int CountContaining(const std::string& substr) { + std::lock_guard lock(mutex_); + int count = 0; + for (const auto& m : messages_) { + if (m.find(substr) != std::string::npos) ++count; + } + return count; + } + + private: + std::vector messages_; + std::mutex mutex_; +}; + +// ------------------------------------------------------------------- +// Manual (non-drain) Move must NOT drop the source client's LOCAL_DISK +// backup. Only drain-originated moves drop the redundant LOCAL_DISK +// metadata (MoveEnd is gated on is_drain). This guards the public +// create_move_task API against silently discarding a durable SSD backup +// (and leaking its file) on a still-live node. +// +// Strategy: manually move every key's MEMORY replica client1 -> client2, +// then drain client1 -> client3. If the manual move preserved client1's +// LOCAL_DISK replicas, the drain migrates them via the LOCAL_DISK path +// (emitting local_disk_transfer_success) and the data becomes readable on +// client3. If the LOCAL_DISK metadata had been wrongly dropped by the +// manual move, the drain would find nothing to migrate, succeeded_units +// would be 0, and client3 would never receive the data. +// ------------------------------------------------------------------- +TEST_F(LocalDiskDrainTest, ManualMovePreservesLocalDiskBackup) { + const std::string client1_addr = "127.0.0.1:17813"; + const std::string client2_addr = "127.0.0.1:17814"; + const std::string client3_addr = "127.0.0.1:17815"; + fs::path ssd1 = tmp_root_ / "c1_ssd"; + fs::path ssd2 = tmp_root_ / "c2_ssd"; + fs::path ssd3 = tmp_root_ / "c3_ssd"; + + SetupOffloadClient(client1_addr, ssd1, client1_); + SetupOffloadClient(client2_addr, ssd2, client2_); + SetupOffloadClient(client3_addr, ssd3, client3_); + ASSERT_EQ(client1_->mountLocalDiskSegment(true), 0) + << "Failed to mount LOCAL_DISK segment on client1"; + + // Phase 1: Put keys on client1 (MEMORY + LOCAL_DISK on client1). + constexpr int kKeyCount = 4; + // Mix of 4096-aligned and deliberately NON-4096-aligned object sizes. + // These keys become LOCAL_DISK-only after the manual move below, so the + // drain migrates them through LoadBatchFromLocalDisk -> + // BucketStorageBackend::BatchLoad, exercising the O_DIRECT offset/size + // correction that a uniform 4096 size would never stress. + const size_t kSizes[kKeyCount] = {4096, 7001, 12000, 5000}; + const size_t kMaxSize = 12000; + std::vector keys, values; + std::vector sizes; + for (int i = 0; i < kKeyCount; ++i) { + std::string key = "manual_move_key_" + std::to_string(i); + // Position- AND key-varying pattern so a byte-for-byte compare detects + // intra-object slice mis-ordering / wrong O_DIRECT offset. + std::string value = MakePatternValue(kSizes[i], i); + std::span span(value.data(), value.size()); + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = client1_addr; + ASSERT_EQ(client1_->put(key, span, config), 0) + << "Put failed for " << key; + keys.push_back(key); + values.push_back(value); + sizes.push_back(kSizes[i]); + } + ASSERT_TRUE(WaitForLocalDiskFiles(ssd1, std::chrono::seconds(120))) + << "LOCAL_DISK replicas were not offloaded to client1"; + + // Phase 2: Manually move each key's MEMORY replica client1 -> client2. + // create_move_task => is_drain=false, so the LOCAL_DISK replica on + // client1 must be preserved by MoveEnd. + for (int i = 0; i < kKeyCount; ++i) { + auto task = + client1_->create_move_task(keys[i], client1_addr, client2_addr); + ASSERT_TRUE(task.has_value()) + << "create_move_task failed for " << keys[i]; + bool done = false; + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(60); + while (std::chrono::steady_clock::now() < deadline) { + auto q = client1_->query_task(task.value()); + if (q.has_value() && (q->status == TaskStatus::SUCCESS || + q->status == TaskStatus::FAILED)) { + ASSERT_EQ(q->status, TaskStatus::SUCCESS) + << "manual move task did not succeed for " << keys[i]; + done = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + ASSERT_TRUE(done) << "manual move task timed out for " << keys[i]; + } + + // Phase 3: Verify each key moved to client2 (MEMORY). + for (int i = 0; i < kKeyCount; ++i) { + auto alloc = client2_->client_buffer_allocator_->allocate(kMaxSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client2_->get_into(keys[i], dst, sizes[i]); + ASSERT_EQ(n, static_cast(sizes[i])) + << "post-manual-move get from client2 failed for " << keys[i]; + ASSERT_EQ(std::string(dst, sizes[i]), values[i]); + } + LOG(INFO) << "[ManualMove] moved " << kKeyCount + << " MEMORY replicas client1 -> client2 (is_drain=false); " + "client1 LOCAL_DISK backups expected to be preserved"; + + // Phase 4: Wait for KV leases to expire so the drain scheduler will not + // block the keys (default TTL 5s). + std::this_thread::sleep_for(std::chrono::seconds(6)); + + // Phase 5: Drain client1 -> client3. Capture logs to prove the + // LOCAL_DISK path ran. + CapturingLogSink sink; + google::AddLogSink(&sink); + HttpQueryDrainJobResponse final_job; + { + CreateDrainJobRequest request; + request.segments = {client1_addr}; + request.target_segments = {client3_addr}; + request.max_concurrency = 1; + auto create_res = CreateDrainJobViaHttp(request); + ASSERT_TRUE(create_res.has_value()); + ASSERT_TRUE(create_res->success) << create_res->error_message; + bool ok = WaitForJobCompletionViaHttp(create_res->job_id, &final_job, + std::chrono::seconds(120)); + google::RemoveLogSink(&sink); + ASSERT_TRUE(ok) << "Drain did not complete (status=" + << (final_job.status_name.empty() + ? "" + : final_job.status_name) + << ")"; + } + EXPECT_EQ(final_job.status_name, "SUCCEEDED"); + EXPECT_EQ(final_job.failed_units, 0u); + // All keys must migrate via the LOCAL_DISK path. If the manual move had + // wrongly dropped the LOCAL_DISK metadata, there would be nothing to + // migrate and succeeded_units would be 0. + EXPECT_EQ(final_job.succeeded_units, static_cast(kKeyCount)) + << "Expected all " << kKeyCount + << " LOCAL_DISK replicas (preserved by the manual move) to be drained"; + + // Direct evidence the LOCAL_DISK drain path ran for every key. + const int local_disk_transfers = + sink.CountContaining("local_disk_transfer_success"); + LOG(INFO) << "[ManualMove] local_disk_transfer_success count=" + << local_disk_transfers << " (expect >= " << kKeyCount << ")"; + EXPECT_GE(local_disk_transfers, kKeyCount) + << "Manual move must preserve LOCAL_DISK; drain should then migrate " + "all " + << kKeyCount << " keys via the LOCAL_DISK path"; + + // Source segment must reach DRAINED (no MEMORY, no LOCAL_DISK left). + { + auto st = QuerySegmentStatusViaHttp(client1_addr); + ASSERT_TRUE(st.has_value()); + EXPECT_EQ(st->status_name, "DRAINED"); + } + + // Phase 6: Data must be byte-for-byte readable from client3 (the drain + // target), proving the preserved LOCAL_DISK data survived migration. + for (int i = 0; i < kKeyCount; ++i) { + auto alloc = client3_->client_buffer_allocator_->allocate(kMaxSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client3_->get_into(keys[i], dst, sizes[i]); + ASSERT_EQ(n, static_cast(sizes[i])) + << "get from drain target client3 failed for " << keys[i]; + ASSERT_EQ(std::string(dst, sizes[i]), values[i]) + << "data mismatch on client3 for " << keys[i]; + } +} + +// ------------------------------------------------------------------- +// Failure path: when a LOCAL_DISK-only replica's backing SSD file is gone at +// drain time, LoadBatchFromLocalDisk fails and the drain unit must fail +// *gracefully*: the job must reach a terminal FAILED state (never hang), the +// source segment must NOT be marked DRAINED, no phantom success is counted, +// and the live (MEMORY) copy of the data must remain intact. This is the ONLY +// test that exercises the LOCAL_DISK transfer error/revoke branch in +// Client::ExecuteReplicaTransfer -- the happy-path tests never reach it. +// +// Setup: put keys on client1 (MEMORY + LOCAL_DISK), manually move the MEMORY +// replica client1 -> client2 (is_drain=false, so client1 keeps its +// LOCAL_DISK-only replica), DELETE client1's backing SSD files to simulate a +// lost/corrupt file, then drain client1 -> client3. +// ------------------------------------------------------------------- +TEST_F(LocalDiskDrainTest, DrainFailsGracefullyWhenLocalDiskDataMissing) { + const std::string client1_addr = "127.0.0.1:17813"; + const std::string client2_addr = "127.0.0.1:17814"; + const std::string client3_addr = "127.0.0.1:17815"; + fs::path ssd1 = tmp_root_ / "c1_ssd"; + fs::path ssd2 = tmp_root_ / "c2_ssd"; + fs::path ssd3 = tmp_root_ / "c3_ssd"; + + SetupOffloadClient(client1_addr, ssd1, client1_); + SetupOffloadClient(client2_addr, ssd2, client2_); + SetupOffloadClient(client3_addr, ssd3, client3_); + ASSERT_EQ(client1_->mountLocalDiskSegment(true), 0); + + constexpr int kKeyCount = 4; + constexpr size_t kValueSize = 4096; + std::vector keys, values; + for (int i = 0; i < kKeyCount; ++i) { + std::string key = "missing_ld_key_" + std::to_string(i); + std::string value = MakePatternValue(kValueSize, 100 + i); + std::span span(value.data(), value.size()); + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = client1_addr; + ASSERT_EQ(client1_->put(key, span, config), 0) << "Put failed " << key; + keys.push_back(key); + values.push_back(value); + } + ASSERT_TRUE(WaitForLocalDiskFiles(ssd1)) << "offload did not materialize"; + + // Move the MEMORY replica off client1 so client1 is LOCAL_DISK-only. + for (int i = 0; i < kKeyCount; ++i) { + auto task = + client1_->create_move_task(keys[i], client1_addr, client2_addr); + ASSERT_TRUE(task.has_value()); + bool done = false; + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(60); + while (std::chrono::steady_clock::now() < deadline) { + auto q = client1_->query_task(task.value()); + if (q.has_value() && (q->status == TaskStatus::SUCCESS || + q->status == TaskStatus::FAILED)) { + ASSERT_EQ(q->status, TaskStatus::SUCCESS); + done = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + ASSERT_TRUE(done) << "manual move timed out for " << keys[i]; + } + + // Inject the fault: delete every backing file under client1's SSD dir. + // client1 is now LOCAL_DISK-only for these keys; the live MEMORY copy is + // on client2. A subsequent LOCAL_DISK load on client1 must now fail at + // open()/read() (BucketStorageBackend::BatchLoad reopens the file fresh). + { + std::error_code ec; + size_t removed = 0; + for (auto it = fs::recursive_directory_iterator(ssd1, ec); + it != fs::recursive_directory_iterator(); ++it) { + if (it->is_regular_file()) { + std::error_code rec; + if (fs::remove(it->path(), rec)) ++removed; + } + } + LOG(INFO) << "[MissingLD] deleted " << removed + << " backing file(s) under " << ssd1.string(); + ASSERT_GT(removed, 0u) << "no SSD files to delete; setup invalid"; + } + + // Let KV leases expire so the scheduler will actually plan the units. + std::this_thread::sleep_for(std::chrono::seconds(6)); + + // Drain client1 -> client3. The LOCAL_DISK load must fail for every key. + CapturingLogSink sink; + google::AddLogSink(&sink); + HttpQueryDrainJobResponse final_job; + bool completed_ok = false; + { + CreateDrainJobRequest request; + request.segments = {client1_addr}; + request.target_segments = {client3_addr}; + request.max_concurrency = 1; + auto create_res = CreateDrainJobViaHttp(request); + ASSERT_TRUE(create_res.has_value()); + ASSERT_TRUE(create_res->success) << create_res->error_message; + // Poll for ANY terminal state; we expect FAILED, not SUCCEEDED. + completed_ok = WaitForJobCompletionViaHttp( + create_res->job_id, &final_job, std::chrono::seconds(150)); + google::RemoveLogSink(&sink); + } + + // 1) The job must reach a TERMINAL state (no hang). final_job.status_name + // is populated only when a terminal state was observed. + ASSERT_FALSE(final_job.status_name.empty()) + << "drain job never reached a terminal state (hang?)"; + // 2) It must be FAILED, not SUCCEEDED: the LOCAL_DISK data is unreadable. + EXPECT_FALSE(completed_ok); + EXPECT_EQ(final_job.status_name, "FAILED"); + EXPECT_EQ(final_job.succeeded_units, 0u) + << "no key could migrate; succeeded_units must be 0"; + EXPECT_GT(final_job.failed_units, 0u); + + // 3) The LOCAL_DISK load-failure branch must have been taken. + EXPECT_GT(sink.CountContaining("local_disk_load_failed"), 0) + << "expected the LOCAL_DISK load-failure error branch to be hit"; + + // 4) The source segment must NOT be DRAINED (it still owns LD replicas + // that could not be migrated). + { + auto st = QuerySegmentStatusViaHttp(client1_addr); + ASSERT_TRUE(st.has_value()); + EXPECT_NE(st->status_name, "DRAINED") + << "segment must not be DRAINED when LOCAL_DISK migration failed"; + } + + // 5) The live copy is intact: client2 still serves the data byte-for-byte, + // proving the failed drain did not corrupt or drop the surviving + // replica. + for (int i = 0; i < kKeyCount; ++i) { + auto alloc = client2_->client_buffer_allocator_->allocate(kValueSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client2_->get_into(keys[i], dst, kValueSize); + ASSERT_EQ(n, static_cast(kValueSize)) + << "live MEMORY copy on client2 lost for " << keys[i]; + ASSERT_EQ(std::string(dst, kValueSize), values[i]) + << "live copy corrupted for " << keys[i]; + } +} + +class LocalDiskDrainEvictTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("LocalDiskDrainEvictTest"); + FLAGS_logtostderr = 1; + } + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + void SetUp() override { + // Reset the process-wide memory metrics so this test starts from a + // clean slate. MasterMetricManager is a singleton; residual capacity + // left over from earlier tests (whose segment unmounts can fail + // during teardown with RPC_FAIL) would otherwise inflate + // mem_total_capacity_ and skew get_global_mem_used_ratio(), which + // would prevent the watermark-based eviction from triggering. + MasterMetricManager::instance().reset_total_mem_capacity(); + MasterMetricManager::instance().reset_allocated_mem_size(); + google::AddLogSink(&log_sink_); + tmp_root_ = fs::temp_directory_path() / + ("mc_ldevict_" + std::to_string(::getpid())); + fs::create_directories(tmp_root_); + ::setenv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "1", 1); + } + + void TearDown() override { + google::RemoveLogSink(&log_sink_); + if (client1_) client1_->tearDownAll(); + if (client2_) client2_->tearDownAll(); + master_.Stop(); + std::error_code ec; + fs::remove_all(tmp_root_, ec); + } + + void SetupClient(const std::string& addr, const fs::path& ssd_path, + std::shared_ptr& out, + size_t seg_size = 64 * 1024 * 1024) { + fs::create_directories(ssd_path); + const std::string rdma_devices = + (FLAGS_protocol == "rdma") ? FLAGS_device_name : std::string(""); + out = RealClient::create(); + ASSERT_NE(out, nullptr); + constexpr size_t kBufSize = 64 * 1024 * 1024; + ASSERT_EQ(out->setup_real(addr, "P2PHANDSHAKE", seg_size, kBufSize, + FLAGS_protocol, rdma_devices, + master_.master_address(), nullptr, "", true, + ssd_path.string(), "default"), + 0) + << "setup_real failed for " << addr; + } + + struct SsdStats { + size_t files{0}; + uint64_t bytes{0}; + }; + SsdStats CountSsd(const fs::path& p) { + SsdStats s; + std::error_code ec; + for (auto it = fs::recursive_directory_iterator(p, ec); + it != fs::recursive_directory_iterator(); ++it) { + if (it->is_regular_file()) { + s.files++; + s.bytes += it->file_size(ec); + } + } + return s; + } + + bool WaitFiles(const fs::path& p, + std::chrono::seconds timeout = std::chrono::seconds(120)) { + auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + std::error_code ec; + for (auto it = fs::recursive_directory_iterator(p, ec); + it != fs::recursive_directory_iterator(); ++it) { + if (it->is_regular_file()) return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + return false; + } + + // Fill client1's 64 MB segment (the ONLY mounted segment) with 4 MB + // objects. Since total capacity = 64 MB, writing 64 MB pushes the + // global mem ratio to 1.0 > watermark (0.95), triggering eviction + // automatically via the eviction thread (10 ms cycle). No "trigger + // Put" needed. + void RunEvictionDrainTest(bool offload_on_evict) { + const std::string src_addr = "127.0.0.1:17813"; + const std::string dst_addr = "127.0.0.1:17814"; + fs::path ssd1 = tmp_root_ / "c1_ssd"; + fs::path ssd2 = tmp_root_ / "c2_ssd"; + + InProcMasterConfig cfg = InProcMasterConfigBuilder() + .set_enable_offload(true) + .set_offload_on_evict(offload_on_evict) + .set_eviction_high_watermark_ratio(0.5) + .build(); + ASSERT_TRUE(master_.Start(cfg)); + + // Only set up client1 initially. Total capacity = 64 MB. + SetupClient(src_addr, ssd1, client1_); + ASSERT_EQ(client1_->mountLocalDiskSegment(true), 0); + + constexpr size_t kValSize = 4 * 1024 * 1024; // 4 MB + constexpr int kKeyCount = 16; // 16 x 4 MB = 64 MB + const std::string src_seg = src_addr; + + // Phase 1: Fill the segment. Ratio reaches 1.0 > 0.95 watermark. + std::vector keys, vals; + for (int i = 0; i < kKeyCount; ++i) { + std::string k = "evict_key_" + std::to_string(i); + std::string v = MakePatternValue(kValSize, i); + std::span sp(v.data(), v.size()); + ReplicateConfig rc; + rc.replica_num = 1; + rc.preferred_segment = src_seg; + ASSERT_EQ(client1_->put(k, sp, rc), 0) << "Put failed for " << k; + keys.push_back(k); + vals.push_back(v); + } + LOG(INFO) << "[Phase 1] offload_on_evict=" << offload_on_evict + << " stored=" << kKeyCount << " x " << kValSize + << " bytes = " << (uint64_t)kKeyCount * kValSize + << " bytes (segment full, ratio=1.0 > watermark 0.95)"; + + // Phase 1.5: Verify every key is readable from client1 right after + // Put. This proves the first leg of EVERY key's route: Put lands the + // data in client1's MEMORY (segment 17813), regardless of whether the + // key is later evicted to LOCAL_DISK or stays in memory. + { + auto alloc = client1_->client_buffer_allocator_->allocate(kValSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + for (int i = 0; i < kKeyCount; ++i) { + auto n = client1_->get_into(keys[i], dst, kValSize); + ASSERT_EQ(n, static_cast(kValSize)) + << "post-Put get_into failed for " << keys[i]; + ASSERT_EQ(std::string(dst, kValSize), vals[i]) + << "post-Put data mismatch for " << keys[i]; + } + LOG(INFO) << "[Phase 1.5] All " << kKeyCount + << " keys verified in client1 MEMORY right after Put " + "(route leg 1: Put -> client1 mem)"; + } + + // Phase 2: For immediate-offload, SSD files exist right away. + if (!offload_on_evict) { + ASSERT_TRUE(WaitFiles(ssd1)) + << "Immediate offload: SSD files not found"; + auto s = CountSsd(ssd1); + LOG(INFO) << "[Phase 2] Immediate offload SSD: files=" << s.files + << " bytes=" << s.bytes; + } + + // Phase 3: Wait for KV leases to expire (TTL = 5 s). The eviction + // thread (10 ms cycle) will detect ratio > 0.5 and evict objects + // whose leases have expired. + std::this_thread::sleep_for(std::chrono::seconds(6)); + + // Phase 4: Wait for the offload to materialise on SSD. + // With offload_on_evict=true, eviction pushes objects to the + // offloading queue; the client writes them to SSD asynchronously. + // With offload_on_evict=false, SSD files already exist (Phase 2). + ASSERT_TRUE(WaitFiles(ssd1, std::chrono::seconds(120))) + << "SSD files not found after eviction"; + { + auto s = CountSsd(ssd1); + LOG(INFO) << "[Phase 4] Post-eviction SSD: files=" << s.files + << " bytes=" << s.bytes; + } + + // Phase 4.5: Wait until eviction has actually FREED the MEMORY + // replicas. CRITICAL: this must happen BEFORE client2 is mounted. + // The eviction watermark is GLOBAL (get_global_mem_used_ratio = + // allocated / capacity summed across ALL mounted segments). With only + // client1's 64 MB segment mounted, the ratio sits at 1.0 while MEMORY + // is pinned for offload; once offload completes, eviction removes the + // MEMORY replicas (the offload_on_evict two-step) and the ratio falls. + // Polling until it drops below the 0.5 watermark proves some keys are + // now LOCAL_DISK-only. If client2 (256 MB) were mounted first, the + // ratio would be diluted to ~0.2 and eviction would stop BEFORE freeing + // MEMORY -- exactly the bug this phase guards against. + double post_evict_ratio = 1.0; + { + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(60); + post_evict_ratio = + MasterMetricManager::instance().get_global_mem_used_ratio(); + while (post_evict_ratio > 0.5 && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + post_evict_ratio = + MasterMetricManager::instance().get_global_mem_used_ratio(); + } + LOG(INFO) + << "[Phase 4.5] Post-eviction global mem ratio=" + << post_evict_ratio + << " (<= 0.5 proves MEMORY replicas were freed; some keys " + "are now LOCAL_DISK-only)"; + ASSERT_LE(post_evict_ratio, 0.5) + << "Eviction did not free MEMORY before timeout; the " + "LOCAL_DISK-only state needed to exercise the LOCAL_DISK " + "drain path was not reached"; + } + + // Phase 5: Now set up client2 as the drain target. Use a LARGE + // segment (256 MB) so that after the drain the global mem ratio + // stays well below the 0.5 watermark and the eviction thread does NOT + // evict the freshly-migrated keys off client2 (they have only a MEMORY + // replica and an expired lease, so a low watermark would otherwise + // evict them -> OBJECT_NOT_FOUND). Mounting it only AFTER Phase 4.5 + // keeps the ratio high during eviction so MEMORY actually gets freed. + SetupClient(dst_addr, ssd2, client2_, /*seg_size=*/256 * 1024 * 1024); + + // Phase 6: Drain source -> target. + CreateDrainJobRequest req; + req.segments = {src_seg}; + req.target_segments = {dst_addr}; + req.max_concurrency = 1; + auto cr = CreateDrainJob(req); + ASSERT_TRUE(cr.has_value()) << "CreateDrainJob HTTP failed"; + ASSERT_TRUE(cr->success) << cr->error_message; + + HttpQueryDrainJobResponse fj; + ASSERT_TRUE(WaitJob(cr->job_id, &fj)) + << "Drain timed out (status=" + << (fj.status_name.empty() ? "" : fj.status_name) << ")"; + EXPECT_EQ(fj.status_name, "SUCCEEDED"); + EXPECT_EQ(fj.failed_units, 0u); + LOG(INFO) << "[Phase 6] Drain " << fj.status_name + << " succeeded=" << fj.succeeded_units + << " failed=" << fj.failed_units + << " blocked=" << fj.blocked_units + << " migrated_bytes=" << fj.migrated_bytes; + + // Phase 7: Every stored key must be byte-for-byte readable on target. + for (int i = 0; i < kKeyCount; ++i) { + auto alloc = client2_->client_buffer_allocator_->allocate(kValSize); + ASSERT_TRUE(alloc.has_value()); + auto* dst = static_cast(alloc->ptr()); + auto n = client2_->get_into(keys[i], dst, kValSize); + ASSERT_EQ(n, static_cast(kValSize)) + << "get_into failed for " << keys[i]; + ASSERT_EQ(std::string(dst, kValSize), vals[i]) + << "Data mismatch for " << keys[i]; + } + LOG(INFO) << "[Phase 7] All " << kKeyCount + << " keys verified on target (offload_on_evict=" + << offload_on_evict << ")"; + + // Route proof -- the LOCAL_DISK drain path WAS exercised: + // Phase 4.5 established the global mem ratio fell to post_evict_ratio + // (<= 0.5). With a 64 MB segment holding 16 x 4 MB keys, ratio <= 0.5 + // means <= 32 MB stayed resident, so at least 8 keys had their MEMORY + // replica FREED and became LOCAL_DISK-only (data only on client1 SSD). + // Phase 6 established succeeded_units == kKeyCount, i.e. ALL 16 keys + // were migrated. A LOCAL_DISK-only key has NO memory replica to move, + // so it can only have travelled the LOCAL_DISK path: client1 SSD -> + // staging buffer (LoadBatchFromLocalDisk) -> TransferWrite -> client2 + // MEMORY. Hence the drain necessarily exercised the LOCAL_DISK path. + const int min_local_disk_only = + kKeyCount - + static_cast(std::ceil(post_evict_ratio * kKeyCount)); + LOG(INFO) << "[Route] MEMORY-key route (x" + << (kKeyCount - min_local_disk_only) + << "): Put->client1 mem -> MEMORY-path TransferWrite -> " + "client2 mem"; + LOG(INFO) << "[Route] LOCAL_DISK-key route (x" << min_local_disk_only + << "): Put->client1 mem -> evict->client1 SSD -> " + "LoadBatchFromLocalDisk->staging buffer -> TransferWrite " + "-> client2 mem"; + EXPECT_GE(min_local_disk_only, 1) + << "Expected at least one LOCAL_DISK-only key to exercise the " + "LOCAL_DISK drain path"; + + // Direct log evidence: each LOCAL_DISK-only key migrated via the + // LOCAL_DISK path emits action=replica_move_local_disk_transfer_success + // (SSD -> staging buffer via LoadBatchFromLocalDisk -> TransferWrite). + // The captured count must reach min_local_disk_only, proving the data + // really passed through LoadBatchFromLocalDisk + TransferWrite (not + // merely inferred from the memory ratio). + const int local_disk_transfers = + log_sink_.CountContaining("local_disk_transfer_success"); + LOG(INFO) << "[Route] local_disk_transfer_success log count=" + << local_disk_transfers + << " (expect >= " << min_local_disk_only << ")"; + EXPECT_GE(local_disk_transfers, min_local_disk_only) + << "Expected the LOCAL_DISK drain path to emit at least " + << min_local_disk_only + << " local_disk_transfer_success logs " + "(SSD->staging buffer->TransferWrite)"; + } + + // ----- HTTP helpers (same as LocalDiskDrainTest) ----- + tl::expected CreateDrainJob( + const CreateDrainJobRequest& request) { + std::string body; + struct_json::to_json(request, body); + auto r = HttpPostJson( + master_.http_metrics_base() + "/api/v1/drain_jobs", body); + if (!r.has_value()) return tl::unexpected(r.error()); + HttpCreateDrainJobResponse p; + struct_json::from_json(p, r.value()); + return p; + } + + tl::expected QueryJob( + const std::string& id) { + auto r = HttpGet(master_.http_metrics_base() + + "/api/v1/drain_jobs/query?job_id=" + id); + if (!r.has_value()) return tl::unexpected(r.error()); + HttpQueryDrainJobResponse p; + struct_json::from_json(p, r.value()); + return p; + } + + bool WaitJob(const std::string& id, HttpQueryDrainJobResponse* out, + std::chrono::seconds timeout = std::chrono::seconds(120)) { + auto start = std::chrono::steady_clock::now(); + while (std::chrono::steady_clock::now() - start < timeout) { + auto qr = QueryJob(id); + if (qr.has_value()) { + if (qr->status_name == "SUCCEEDED" || + qr->status_name == "FAILED" || + qr->status_name == "CANCELED") { + if (out) *out = qr.value(); + return qr->status_name == "SUCCEEDED"; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + return false; + } + + InProcMaster master_; + std::shared_ptr client1_; + std::shared_ptr client2_; + fs::path tmp_root_; + CapturingLogSink log_sink_; +}; + +// Scenario 1: offload_on_evict=true 鈥?offload deferred to eviction time. +// Objects are NOT offloaded at PutEnd; they are offloaded only when the +// eviction thread selects them. After eviction + offload, some keys have +// only LOCAL_DISK replicas. The drain must migrate them via the LOCAL_DISK +// path (BatchLoad -> TransferWrite). +TEST_F(LocalDiskDrainEvictTest, DrainAfterEviction_OffloadOnEvict) { + RunEvictionDrainTest(/*offload_on_evict=*/true); +} + +// Scenario 2: offload_on_evict=false 鈥?offload immediate at PutEnd. +// Objects are offloaded to SSD right after Put. When eviction fires, +// keys already have LOCAL_DISK replicas, so eviction just removes the +// MEMORY replica. The drain handles a mix of MEMORY-only and +// LOCAL_DISK-only keys. +TEST_F(LocalDiskDrainEvictTest, DrainAfterEviction_ImmediateOffload) { + RunEvictionDrainTest(/*offload_on_evict=*/false); +} + +} // namespace testing +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, false /* remove_flags */); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-store/tests/test_server_helpers.h b/mooncake-store/tests/test_server_helpers.h index 765eb4429c..b6c763d8ec 100644 --- a/mooncake-store/tests/test_server_helpers.h +++ b/mooncake-store/tests/test_server_helpers.h @@ -136,6 +136,9 @@ class InProcMaster { if (config.quota_bytes.has_value()) { wms_cfg.quota_bytes = config.quota_bytes.value(); } + wms_cfg.offload_on_evict = config.offload_on_evict.has_value() + ? config.offload_on_evict.value() + : false; wms_cfg.enable_cxl = config.enable_cxl.has_value() ? config.enable_cxl.value()