[Bugfix][Store] Register ShmHelper memory with SPDK for NoF zero-copy transfers

docs/source/deployment/mooncake-store-deployment-guide.md

@@ -1281,6 +1281,7 @@ Do not run binaries from before and after checksum support was introduced in the
 | `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `512MB`, `1GB` |
 | `MC_MMAP_ARENA_POOL_SIZE` | unset | Pre-allocated arena pool size (e.g., `8gb`). Explicitly set to enable the arena |
 | `MC_DISABLE_MMAP_ARENA` | unset | Disable arena, fall back to per-call `mmap()`. Accepts `1`/`true`/`yes`/`on` (or `0`/`false`/`no`/`off`) |
+| `MC_STORE_REGISTER_SPDK` | unset | Set `1` to register `ShmHelper`-allocated shared memory (host pool, dummy local buffer) with SPDK for NoF zero-copy transfers. Forces HugeTLB backing for those allocations, defaulting to 2MB hugepages when `MC_STORE_USE_HUGEPAGE` is unset. Must be set on BOTH the dummy and the real process (SPDK registration is per-process) |
 
 RDMA Store segments backed by HugeTLB are populated in parallel immediately
 before transfer-engine registration. No additional population-mode setting is
@@ -1297,6 +1298,25 @@ with its `mbind()` region before touching pages. The mmap arena retains its
 eager `MAP_POPULATE` behavior for DMA safety; set `MC_DISABLE_MMAP_ARENA=1` if
 the deferred direct-mmap path is desired while the arena is otherwise enabled.
 
+For NoF (NVMe-oF) zero-copy, `MC_STORE_REGISTER_SPDK=1` registers the shared
+memory allocated by `ShmHelper` (SGLang host pool, dummy local buffer) with
+SPDK (`spdk_mem_register`) so the NoF RDMA transport can DMA to/from it
+directly — without it those buffers fail with `No translation for ptr`.
+`spdk_mem_register` is per-process, so set this switch on BOTH the dummy
+(sender) and the real client (receiver): the dummy registers its own mapping
+in `ShmHelper::allocate`, and the real client registers its separate mapping
+of the same shared fd in `RealClient::map_shm_internal_with_device`. Setting
+it on only one process leaves the other without an SPDK translation and NoF
+transfers still fail with `No translation for ptr`. SPDK
+registration in iova=pa mode requires PHYSICALLY 2MB-aligned memory, which only
+HugeTLB pages satisfy, so this switch forces HugeTLB for the affected
+allocations even when `MC_STORE_USE_HUGEPAGE` is unset; it then defaults to 2MB
+hugepages (set `MC_STORE_USE_HUGEPAGE=1` and `MC_STORE_HUGEPAGE_SIZE=1GB` for
+1GB). Reserve enough HugeTLB pages (`/proc/sys/vm/nr_hugepages`) for the host
+pool plus any hugepage-backed segments; when the pool is exhausted the first
+allocation aborts with a clear error naming the hugepage size and count needed
+rather than silently degrading.
+
 #### yalantinglibs Log Level
 
 ```bash

mooncake-store/include/common/mmap_aligned.h

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <sys/mman.h>
+
+#include "common/client_buffer_allocation.h"
+
+namespace mooncake {
+
+/**
+ * @brief Place a memfd mapping at a 2MB-aligned address while keeping the file
+ * offset at 0, so the returned base matches what a peer process mapping the fd
+ * at offset 0 sees (cross-process fd sharing is unchanged).
+ *
+ * SPDK before 26.09 (e.g. the v23.01.1 pinned by dependencies.sh) rejects
+ * spdk_mem_register() unless both the start and the length are 2MB-aligned;
+ * plain mmap only guarantees 4KB alignment. Reserve a PROT_NONE region large
+ * enough to contain a 2MB-aligned window, then MAP_FIXED the memfd into it
+ * (which only replaces the reservation, never live mappings), and release the
+ * unaligned head/tail. `size` should be a multiple of 2MB (callers align it).
+ *
+ * NOTE: virtual 2MB alignment is NECESSARY but NOT sufficient on the pinned
+ * SPDK: in iova=pa mode spdk_mem_register additionally checks each segment's
+ * PHYSICAL address (vtophys pagemap path, `paddr & MASK_2MB`), so the memory
+ * must also be hugepage-backed or registration fails with -EINVAL. HugeTLB
+ * mappings are automatically 2MB/1GB-aligned, so this helper is only needed
+ * where the source cannot be hugepage-backed; today it is used solely by the
+ * receiver (RealClient::map_shm_internal_with_device) for the shared fd
+ * (1GB-hugepage fds whose MAP_FIXED here fails fall back to a plain mmap). The
+ * sender (ShmHelper::allocate) always maps hugepages under
+ * MC_STORE_REGISTER_SPDK=1 and needs no aligned base. SPDK >= 26.09 (commit
+ * d6dc356ff) relaxes the API check to 4KB, but iova=pa still needs hugepages,
+ * so a plain mmap remains insufficient for SPDK registration regardless of
+ * version.
+ *
+ * @param size Mapping length; must be a multiple of 2MB for a clean head/tail.
+ * @param fd memfd to map at offset 0.
+ * @return The 2MB-aligned base, or MAP_FAILED on error.
+ */
+inline void* mmap_shm_2mb_aligned(size_t size, int fd) {
+    const size_t reserve_size = size + SZ_2MB;
+    void* hint = mmap(nullptr, reserve_size, PROT_NONE,
+                      MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    if (hint == MAP_FAILED) {
+        return MAP_FAILED;
+    }
+    const uintptr_t hint_addr = reinterpret_cast<uintptr_t>(hint);
+    const uintptr_t target = static_cast<uintptr_t>(
+        align_up(static_cast<size_t>(hint_addr), SZ_2MB));
+    void* base =
+        mmap(reinterpret_cast<void*>(target), size, PROT_READ | PROT_WRITE,
+             MAP_SHARED | MAP_FIXED | MAP_POPULATE, fd, 0);
+    if (base == MAP_FAILED) {
+        munmap(hint, reserve_size);
+        return MAP_FAILED;
+    }
+    // Release the unaligned head and tail of the reservation.
+    const uintptr_t base_addr = reinterpret_cast<uintptr_t>(base);
+    if (base_addr > hint_addr) {
+        munmap(hint, base_addr - hint_addr);
+    }
+    const uintptr_t reserve_end = hint_addr + reserve_size;
+    const uintptr_t tail = base_addr + size;
+    if (tail < reserve_end) {
+        munmap(reinterpret_cast<void*>(tail), reserve_end - tail);
+    }
+    return base;
+}
+
+}  // namespace mooncake

mooncake-store/include/real_client.h

@@ -992,6 +992,17 @@ class RealClient : public PyClient {
         void *shm_buffer = nullptr;
         size_t shm_size = 0;
         uintptr_t dummy_base_addr = 0;
+        // Whether this receiver-side mapping was registered with SPDK for NoF
+        // zero-copy (see RealClient::map_shm_internal_with_device). Must be
+        // unregistered before munmap on every teardown path.
+        bool spdk_registered = false;
+        // Set when the transfer engine's unregisterLocalMemory() failed during
+        // teardown: TE still holds this region (RdmaTransport does not
+        // deregister the MR when the metadata update fails, and
+        // TransferEngineImpl keeps its region record), so the mapping must not
+        // be munmapped until a later retry releases it (see
+        // RetryQuarantinedShmsLocked).
+        bool te_unregister_pending = false;
         bool is_ascend = false;
         bool is_ipc = false;
         // Ascend physical device id from dummy (dummy-real RPC).
@@ -1129,6 +1140,18 @@ class RealClient : public PyClient {
     void ReleaseAllMountedSegmentRecords();
     void ReleaseAllocatedSegmentRecord(const std::string &segment_id);
     void ReleaseAllAllocatedSegmentRecords();
+
+    // MappedShm whose spdk_mem_unregister() (spdk_registered) or transfer
+    // engine unregisterLocalMemory() (te_unregister_pending) failed during
+    // teardown. The mapping is deliberately retained (never munmapped) so
+    // neither SPDK nor the transfer engine keeps a registration to a
+    // freed/reused VA; retried on later teardown entry points. Guarded by
+    // dummy_client_mutex_.
+    std::vector<MappedShm> quarantine_shms_;
+
+    // Re-attempt TE + SPDK unregister and munmap of quarantine_shms_. Caller
+    // must hold dummy_client_mutex_. Failures stay quarantined.
+    void RetryQuarantinedShmsLocked();
 };
 
 }  // namespace mooncake

mooncake-store/include/shm_helper.h

@@ -23,9 +23,18 @@ class ShmHelper {
     struct ShmSegment {
         int fd = -1;
         void *base_addr = nullptr;
+        // Size of the actual mapping; may be padded up to the hugepage/2MB
+        // boundary when SPDK registration is enabled
+        // (MC_STORE_REGISTER_SPDK=1).
         size_t size = 0;
+        // Size the caller requested in allocate(); == size unless padded.
+        // Callers that must match the original request (e.g.
+        // DummyClient::register_buffer) compare against this, not size, so they
+        // don't re-derive the alignment.
+        size_t requested_size = 0;
         std::string name;
         bool registered = false;
+        bool spdk_registered = false;
         bool is_local = false;
     };
 
@@ -45,6 +54,12 @@ class ShmHelper {
 
     bool is_hugepage() const { return use_hugepage_; }
 
+    // Whether the MC_STORE_REGISTER_SPDK feature is enabled (env == "1").
+    // Single source of truth for the env semantics, shared by ShmHelper's
+    // constructor and the RealClient-side (receiver) registration so both sides
+    // of the dummy/real split gate identically.
+    static bool is_register_spdk_enabled();
+
     ShmHelper(const ShmHelper &) = delete;
     ShmHelper &operator=(const ShmHelper &) = delete;
 
@@ -55,6 +70,7 @@ class ShmHelper {
     std::vector<std::shared_ptr<ShmSegment>> shms_;
     static std::mutex shm_mutex_;
     bool use_hugepage_ = false;
+    bool register_spdk_ = false;
 };
 
 }  // namespace mooncake

mooncake-store/include/spdk/spdk_wrapper.h

@@ -53,6 +53,52 @@ class SpdkWrapper {
     bool ProbeNofSegment(const std::string &tr_str, uint32_t timeout_ms,
                          std::string *error_reason = nullptr);
 
+    /** @brief Register external memory with SPDK for NoF zero-copy transfers.
+     *
+     * Memory returned by Alloc() is already registered with SPDK; use this to
+     * register memory allocated outside of SPDK (e.g. mmap'd shared memory) so
+     * that NoF RDMA transfers can DMA to/from it directly
+     * (spdk_rdma_get_translation).
+     *
+     * The pinned SPDK (v23.01.1) requires BOTH addr and size to be 2MB-aligned
+     * (MASK_2MB), and in iova=pa mode (the typical libibverbs NoF setup) the
+     * PHYSICAL pages must be 2MB-aligned too, which only HugeTLB pages satisfy.
+     * 4KB page-aligned memory always fails with -EINVAL. Callers must pass
+     * 2MB-aligned, hugepage-backed memory (ShmHelper forces this under
+     * MC_STORE_REGISTER_SPDK=1) and check IsRegistrableRange() first: a failure
+     * past that check may leave SPDK's g_mem_reg_map marked, so callers must
+     * attempt UnregisterMemory() and keep the mapping unless it returns 0.
+     *
+     * @param addr Start of the region; must be 2MB-aligned (hugepage-backed).
+     * @param size Region length; must be a multiple of 2MB.
+     * @return 0 on success, non-zero on failure.
+     */
+    int RegisterMemory(void *addr, size_t size);
+
+    /** @brief Unregister memory previously registered via RegisterMemory().
+     *
+     * Only a 0 return proves the range is clean. A non-zero return does NOT
+     * prove that nothing was ever registered: -EINVAL is also what SPDK
+     * returns for a half-registered range, so callers must keep the mapping
+     * alive (and retry) unless this returns 0. See the definition.
+     *
+     * @return 0 on success, non-zero on failure.
+     */
+    int UnregisterMemory(void *addr, size_t size);
+
+    /** @brief Whether RegisterMemory() can be attempted for this range at all.
+     *
+     * spdk_mem_register() rejects a range it cannot represent or that is not
+     * 2MB-aligned before it has marked anything (SPDK v23.01.1,
+     * lib/env_dpdk/memory.c:339-348), so a caller that skips the registration
+     * for such a range still munmaps normally. Any failure past this check may
+     * have left translation state behind and must be treated as such, with one
+     * exception the caller handles separately: -EBUSY means another owner
+     * already holds the range, so nothing new was marked and rolling back would
+     * clear their registration instead.
+     */
+    static bool IsRegistrableRange(void *addr, size_t size);
+
    private:
     struct ProbeBuffer {
         void *ptr{nullptr};

mooncake-store/src/dummy_client.cpp

@@ -956,7 +956,17 @@ int DummyClient::unregister_device_buffer_for_reconnect(void* buffer) {
 
     auto ret = invoke_rpc<&RealClient::unregister_shm_buffer_internal, void>(
         static_cast<uint64_t>(buffer_addr), client_id_);
-    if (ret.has_value()) registered_device_buffers_.erase(buffer_addr);
+    // Same contract as unregister_buffer: the receiver removes the mapping on
+    // INTERNAL_ERROR too (unmapped, or quarantined when the SPDK or transfer
+    // engine unregister fails), and INVALID_PARAMS means the server has no such
+    // mapping, so drop the local bookkeeping then; keep it on RPC_FAIL /
+    // RPC_TIMEOUT where the server may not have run (see unregister_buffer).
+    // No extra locking here: the function-scope lock above is already held.
+    if (ret.has_value() ||
+        (!ret.has_value() && (ret.error() == ErrorCode::INTERNAL_ERROR ||
+                              ret.error() == ErrorCode::INVALID_PARAMS))) {
+        registered_device_buffers_.erase(buffer_addr);
+    }
     return to_py_ret(ret);
 }
 
@@ -1045,14 +1055,20 @@ int DummyClient::register_buffer(void* buffer, size_t size) {
         }
         size = align_up(size, alignment);
     }
-    // Check bounds
+    // Check bounds. The buffer must be the segment base and the size must match
+    // either the caller's original request (ShmSegment::requested_size) or the
+    // padded mapping size (shm->size, aligned up to the hugepage/2MB boundary
+    // for SPDK registration). Both are valid: IPC always sends shm->size, and
+    // legacy callers may pass the padded size directly.
     if (reinterpret_cast<uint8_t*>(buffer) !=
             reinterpret_cast<uint8_t*>(shm->base_addr) ||
-        size != shm->size) {
+        (size != shm->requested_size && size != shm->size)) {
         LOG(ERROR) << "Invalid buffer address or size for registration: "
                       "Buffer addr: "
                    << buffer << ", need addr: " << shm->base_addr
-                   << ", buffer size: " << size << ", need size: " << shm->size;
+                   << ", buffer size: " << size
+                   << ", need size: " << shm->requested_size
+                   << " (padded: " << shm->size << ")";
         return -1;
     }
 
@@ -1126,7 +1142,20 @@ int DummyClient::unregister_buffer(void* buffer) {
     }
     auto ret = invoke_rpc<&RealClient::unregister_shm_buffer_internal, void>(
         reinterpret_cast<uint64_t>(buffer), client_id_);
-    if (ret.has_value()) {
+    // INTERNAL_ERROR means the receiver still executed the teardown and removed
+    // the segment from its active mappings: it unmaps normally, or -- when the
+    // SPDK or transfer engine unregister fails -- quarantines the mapping
+    // (retained, never reused for new registrations). Either way the old
+    // mapping is no longer usable on the receiver, so drop the local flag for
+    // INTERNAL_ERROR and for INVALID_PARAMS (server has no such mapping). Keep
+    // it on RPC_FAIL / RPC_TIMEOUT: the RPC may never have reached the server,
+    // or the receiver may have refused before touching anything (Ascend context
+    // setup failure in unregister_shm_buffer_internal returns RPC_FAIL for
+    // exactly that reason), so the mapping can still be alive and a retry must
+    // reach the receiver.
+    if (ret.has_value() ||
+        (!ret.has_value() && (ret.error() == ErrorCode::INTERNAL_ERROR ||
+                              ret.error() == ErrorCode::INVALID_PARAMS))) {
         shm->registered = false;
     }
     return to_py_ret(ret);