[TransferEngine] Add optional HugeTLB path for create_shared_segment

mooncake-integration/shared_segment.py

@@ -61,6 +61,7 @@ def create_shared_segment(
     comm_group: Any = None,
     mmap: bool = True,
     host_register: bool = False,
+    hugetlb: bool = False,
     *,
     tp_group: Any = None,
 ) -> SharedSegment:
@@ -75,8 +76,10 @@ def create_shared_segment(
 
     ``mmap`` (default True) uses an unnamed memfd and asks the kernel for THP
     (``MADV_HUGEPAGE``) so the span can be 2MiB pages without a HugeTLB pool;
-    4KiB pages if THP cannot allocate. Ranks share those pages through the
-    memfd. Set ``mmap=False`` for the platform VMM fabric path.
+    4KiB pages if THP cannot allocate. Set ``hugetlb=True`` to use
+    ``MFD_HUGETLB`` instead; that consumes ``vm.nr_hugepages`` and fails if
+    the pool cannot back the span. Ranks share those pages through the memfd.
+    Set ``mmap=False`` for the platform VMM fabric path.
     ``host_register`` (default False) HostRegister's mmap pages for
     ``device_id`` so ``tensors().data_ptr()`` is a device VA suitable for TE
     ``location=\"npu\"`` ROCE D2rH; requires ``mmap=True``. Ascend VMM
@@ -87,7 +90,11 @@ def create_shared_segment(
     comm_group = _select_comm_group(comm_group, tp_group)
     if host_register and not mmap:
         raise SharedSegmentError("host_register requires mmap=True")
-    if not shared_segment_supported(mmap=mmap, host_register=host_register):
+    if hugetlb and not mmap:
+        raise SharedSegmentError("hugetlb requires mmap=True")
+    if not shared_segment_supported(
+        mmap=mmap, host_register=host_register, hugetlb=hugetlb
+    ):
         if not mmap:
             raise SharedSegmentError(
                 "This mooncake build has no VMM backend for shared segments"
@@ -96,6 +103,10 @@ def create_shared_segment(
             raise SharedSegmentError(
                 "This mooncake build cannot HostRegister shared-segment pages"
             )
+        if hugetlb:
+            raise SharedSegmentError(
+                "This mooncake build cannot allocate HugeTLB shared-segment pages"
+            )
         raise SharedSegmentError(
             "This mooncake build has no mmap shared-segment backend"
         )
@@ -126,6 +137,7 @@ def create_shared_segment(
             device,
             mmap,
             host_register,
+            hugetlb,
         )
     except RuntimeError as exc:
         create_error = str(exc)
@@ -214,8 +226,10 @@ def total_size(self) -> int:
         return self._segment.size()
 
 
-def shared_segment_supported(mmap: bool = True, host_register: bool = False) -> bool:
-    return _CppSharedSegment.supported(mmap, host_register)
+def shared_segment_supported(
+    mmap: bool = True, host_register: bool = False, hugetlb: bool = False
+) -> bool:
+    return _CppSharedSegment.supported(mmap, host_register, hugetlb)
 
 
 @dataclass(frozen=True)

mooncake-integration/transfer_engine/shared_segment_py.cpp

@@ -36,14 +36,15 @@ void ThrowIfError(const Status& status) {
 std::pair<std::shared_ptr<SharedSegment>, py::bytes> SegmentCreate(
     const std::string& name, uint64_t size, uint32_t world_size,
     uint32_t rank_id, uint32_t owner_rank, int32_t device_id, bool mmap,
-    bool host_register) {
+    bool host_register, bool hugetlb) {
     SharedSegmentOptions options;
     options.world_size = world_size;
     options.rank_id = rank_id;
     options.owner_rank = owner_rank;
     options.device_id = device_id;
     options.mmap = mmap;
     options.host_register = host_register;
+    options.hugetlb = hugetlb;
 
     std::string blob;
     std::shared_ptr<SharedSegment> segment;
@@ -74,9 +75,11 @@ void bind_shared_segment(py::module_& m) {
                     py::arg("world_size"), py::arg("rank_id"),
                     py::arg("owner_rank") = 0, py::arg("device_id") = 0,
                     py::arg("mmap") = true, py::arg("host_register") = false,
+                    py::arg("hugetlb") = false,
                     "Phase one; returns (segment, blob).")
         .def_static("supported", &SharedSegment::Supported,
                     py::arg("mmap") = true, py::arg("host_register") = false,
+                    py::arg("hugetlb") = false,
                     "Whether this build can share memory across processes.")
         .def("complete", &SegmentComplete, py::arg("blobs"),
              "Phase two; blobs are indexed by rank.")

mooncake-transfer-engine/include/shared_segment/shared_segment.h

@@ -58,12 +58,17 @@ struct SharedSegmentOptions {
     int32_t device_id = 0;
     // true: unnamed memfd + mmap. Prefer THP (MADV_HUGEPAGE) so the kernel
     // can back the span with 2MiB pages; 4KiB pages if THP cannot allocate.
-    // Does not use the HugeTLB pool. Ranks share the pages through the memfd.
-    // false: platform VMM with a fabric shareable handle.
+    // Does not use the HugeTLB pool unless `hugetlb` is also true. Ranks share
+    // the pages through the memfd. false: platform VMM with a fabric shareable
+    // handle.
     bool mmap = true;
     // When mmap is true: HostRegister the pages for device_id (Ascend).
     // Needed for TE location=npu ROCE D2rH; off by default.
     bool host_register = false;
+    // When mmap is true: unnamed memfd with MFD_HUGETLB. Requires a HugeTLB
+    // pool (`vm.nr_hugepages`); Create fails if the pool cannot back the span.
+    // Default false keeps the THP path.
+    bool hugetlb = false;
 };
 
 class SharedSegmentBackend;
@@ -84,8 +89,10 @@ class SharedSegment {
 
     // Whether this build can share memory across processes. mmap uses an
     // anonymous memfd (always available). host_register additionally needs
-    // Ascend. VMM (mmap=false) needs a platform fabric backend.
-    static bool Supported(bool mmap = true, bool host_register = false);
+    // Ascend. hugetlb additionally needs mmap. VMM (mmap=false) needs a
+    // platform fabric backend. Runtime HugeTLB pool capacity is not probed.
+    static bool Supported(bool mmap = true, bool host_register = false,
+                          bool hugetlb = false);
 
     // Phase two. `blobs` is indexed by rank. After success, `ready()` is true
     // and `base_addr()` is valid.

mooncake-transfer-engine/src/shared_segment/mmap_shared_segment.cpp

@@ -43,16 +43,35 @@
 #ifndef MFD_CLOEXEC
 #define MFD_CLOEXEC 0x0001U
 #endif
+#ifndef MFD_HUGETLB
+#define MFD_HUGETLB 0x0004U
+#endif
+#ifndef MFD_HUGE_SHIFT
+#define MFD_HUGE_SHIFT 26
+#endif
+#ifndef MFD_HUGE_2MB
+#define MFD_HUGE_2MB (21U << MFD_HUGE_SHIFT)
+#endif
+#ifndef MFD_HUGE_1GB
+#define MFD_HUGE_1GB (30U << MFD_HUGE_SHIFT)
+#endif
+#ifndef MFD_HUGE_512MB
+#define MFD_HUGE_512MB (29U << MFD_HUGE_SHIFT)
+#endif
 
 namespace mooncake {
 namespace {
 constexpr uint16_t kMmapBackendId = 3;
 // Matches the common PMD THP size when sysfs is unavailable.
 constexpr uint64_t kFallbackHugePageSize = 2ULL * 1024 * 1024;
+constexpr uint64_t kTwoMiB = 2ULL * 1024 * 1024;
+constexpr uint64_t kFiveTwelveMiB = 512ULL * 1024 * 1024;
+constexpr uint64_t kOneGiB = 1024ULL * 1024 * 1024;
 constexpr const char* kThpPmdSizePath =
     "/sys/kernel/mm/transparent_hugepage/hpage_pmd_size";
 constexpr const char* kShmemThpEnabledPath =
     "/sys/kernel/mm/transparent_hugepage/shmem_enabled";
+constexpr const char* kMemInfoPath = "/proc/meminfo";
 
 // memfd without MFD_HUGETLB is a tmpfs/shmem inode, so shmem THP applies.
 // Sysfs looks like "always within_size [advise] never deny force".
@@ -98,6 +117,46 @@ uint64_t HugePageSize() {
     return kSize;
 }
 
+// Default HugeTLB page size from /proc/meminfo, not the THP PMD size. On
+// arm64 with 64KiB base pages those two often differ (2MiB vs 512MiB).
+uint64_t HugeTlbPageSize() {
+    static const uint64_t kSize = []() -> uint64_t {
+        FILE* fp = std::fopen(kMemInfoPath, "r");
+        if (fp != nullptr) {
+            char line[256];
+            while (std::fgets(line, sizeof(line), fp) != nullptr) {
+                unsigned long long kib = 0;
+                if (std::sscanf(line, "Hugepagesize: %llu kB", &kib) == 1 &&
+                    kib > 0) {
+                    std::fclose(fp);
+                    return static_cast<uint64_t>(kib) * 1024ULL;
+                }
+            }
+            std::fclose(fp);
+        }
+        return kFallbackHugePageSize;
+    }();
+    return kSize;
+}
+
+unsigned int HugeTlbMemFdSizeFlag() {
+    const uint64_t size = HugeTlbPageSize();
+    if (size == kTwoMiB) {
+        return MFD_HUGE_2MB;
+    }
+    if (size == kFiveTwelveMiB) {
+        return MFD_HUGE_512MB;
+    }
+    if (size == kOneGiB) {
+        return MFD_HUGE_1GB;
+    }
+    return 0;
+}
+
+uint64_t MappingAlign(const SharedSegmentOptions& options) {
+    return options.hugetlb ? HugeTlbPageSize() : HugePageSize();
+}
+
 void AdviseTransparentHugePages(void* addr, uint64_t size) {
 #ifdef MADV_HUGEPAGE
     if (addr == nullptr || size == 0) {
@@ -128,10 +187,21 @@ void AdviseDontFork(void* addr, uint64_t size) {
 #endif
 }
 
-int CreateAnonymousMemFd() {
+int CreateAnonymousMemFd(bool hugetlb) {
 #ifdef __NR_memfd_create
-    return static_cast<int>(syscall(__NR_memfd_create, "mcss", MFD_CLOEXEC));
+    unsigned int flags = MFD_CLOEXEC;
+    if (!hugetlb) {
+        return static_cast<int>(syscall(__NR_memfd_create, "mcss", flags));
+    }
+    const unsigned int sized = flags | MFD_HUGETLB | HugeTlbMemFdSizeFlag();
+    int fd = static_cast<int>(syscall(__NR_memfd_create, "mcss", sized));
+    if (fd < 0 && errno == EINVAL && HugeTlbMemFdSizeFlag() != 0) {
+        fd = static_cast<int>(
+            syscall(__NR_memfd_create, "mcss", flags | MFD_HUGETLB));
+    }
+    return fd;
 #else
+    (void)hugetlb;
     errno = ENOSYS;
     return -1;
 #endif
@@ -177,11 +247,14 @@ void PrefaultAnonymousRange(void* addr, uint64_t size, uint64_t stride) {
     bytes[size - 1] = 0;
 }
 
-Status MapMemFd(int fd, uint64_t size, void* hint, void*& addr) {
+Status MapMemFd(int fd, uint64_t size, void* hint, void*& addr, bool populate) {
     int map_flags = MAP_SHARED;
     if (hint != nullptr) {
         map_flags |= MAP_FIXED;
     }
+    if (populate) {
+        map_flags |= MAP_POPULATE;
+    }
     void* mapped = mmap(hint, size, PROT_READ | PROT_WRITE, map_flags, fd, 0);
     if (mapped == MAP_FAILED) {
         return Status::Memory(std::string("mmap failed: ") +
@@ -191,38 +264,57 @@ Status MapMemFd(int fd, uint64_t size, void* hint, void*& addr) {
     return Status::OK();
 }
 
-// Unnamed memfd (no MFD_HUGETLB, no nr_hugepages). MADV_HUGEPAGE is applied
-// before any page is faulted so the kernel can use THP; otherwise 4KiB pages.
-Status CreateAndMapMemFd(uint64_t size, int& fd, void*& addr) {
-    fd = CreateAnonymousMemFd();
+// mmap=true, hugetlb=false: unnamed memfd (no MFD_HUGETLB, no nr_hugepages).
+// MADV_HUGEPAGE is applied before any page is faulted so the kernel can use
+// THP; otherwise 4KiB pages.
+// mmap=true, hugetlb=true: unnamed memfd with MFD_HUGETLB. MAP_POPULATE
+// consumes the HugeTLB pool immediately and fails if it cannot.
+Status CreateAndMapMemFd(uint64_t size, bool hugetlb, int& fd, void*& addr) {
+    fd = CreateAnonymousMemFd(hugetlb);
     if (fd < 0) {
-        return Status::Memory(std::string("memfd_create failed: ") +
-                              std::strerror(errno));
+        std::string msg = "memfd_create failed: ";
+        msg += std::strerror(errno);
+        if (hugetlb) {
+            msg += " (HugeTLB; check /proc/sys/vm/nr_hugepages)";
+        }
+        return Status::Memory(msg);
     }
     if (ftruncate(fd, static_cast<off_t>(size)) != 0) {
         const int err = errno;
         close(fd);
         fd = -1;
-        return Status::Memory(std::string("ftruncate failed: ") +
-                              std::strerror(err));
+        std::string msg = "ftruncate failed: ";
+        msg += std::strerror(err);
+        if (hugetlb) {
+            msg += " (HugeTLB size must be a multiple of Hugepagesize)";
+        }
+        return Status::Memory(msg);
     }
-    void* hint = ReserveAlignedWindow(size, HugePageSize());
+    const uint64_t align = hugetlb ? HugeTlbPageSize() : HugePageSize();
+    void* hint = ReserveAlignedWindow(size, align);
     if (hint == MAP_FAILED) {
         hint = nullptr;
     }
-    auto status = MapMemFd(fd, size, hint, addr);
+    auto status = MapMemFd(fd, size, hint, addr, /*populate=*/hugetlb);
     if (!status.ok() && hint != nullptr) {
         (void)munmap(hint, size);
         hint = nullptr;
-        status = MapMemFd(fd, size, nullptr, addr);
+        status = MapMemFd(fd, size, nullptr, addr, /*populate=*/hugetlb);
     }
     if (!status.ok()) {
         close(fd);
         fd = -1;
+        if (hugetlb) {
+            return Status::Memory(
+                std::string(status.message()) +
+                " (HugeTLB; check /proc/sys/vm/nr_hugepages)");
+        }
         return status;
     }
-    AdviseTransparentHugePages(addr, size);
-    PrefaultAnonymousRange(addr, size, HugePageSize());
+    if (!hugetlb) {
+        AdviseTransparentHugePages(addr, size);
+        PrefaultAnonymousRange(addr, size, HugePageSize());
+    }
     AdviseDontFork(addr, size);
     return Status::OK();
 }
@@ -288,11 +380,10 @@ class MmapSharedSegmentBackend : public SharedSegmentBackend {
    public:
     ~MmapSharedSegmentBackend() override { Release(); }
 
-    uint64_t Granularity(
-        const SharedSegmentOptions& /*options*/) const override {
-        // Align length to the THP PMD size so MADV_HUGEPAGE can use 2MiB
-        // pages. 4KiB fallback still accepts a PMD-aligned length.
-        return HugePageSize();
+    uint64_t Granularity(const SharedSegmentOptions& options) const override {
+        // HugeTLB must match Hugepagesize. THP aligns to the PMD size so
+        // MADV_HUGEPAGE can use 2MiB pages; 4KiB fallback still accepts it.
+        return MappingAlign(options);
     }
 
     Status CreateOwner(uint64_t size, const SharedSegmentOptions& options,
@@ -341,13 +432,18 @@ Status MmapSharedSegmentBackend::CreateOwner(
     std::vector<uint8_t>& handle) {
     int fd = -1;
     void* mapped = nullptr;
-    auto status = CreateAndMapMemFd(size, fd, mapped);
+    auto status = CreateAndMapMemFd(size, options.hugetlb, fd, mapped);
     if (!status.ok()) {
         return status;
     }
-    LOG(INFO) << "Shared segment mmap: unnamed memfd + THP hint (no HugeTLB), "
-              << "size " << size
-              << ", shmem_thp=" << (ShmemThpEnabled() ? "on" : "off");
+    if (options.hugetlb) {
+        LOG(INFO) << "Shared segment mmap: unnamed memfd + HugeTLB, size "
+                  << size << ", hugepage=" << HugeTlbPageSize();
+    } else {
+        LOG(INFO) << "Shared segment mmap: unnamed memfd + THP hint "
+                  << "(no HugeTLB), size " << size
+                  << ", shmem_thp=" << (ShmemThpEnabled() ? "on" : "off");
+    }
 
     addr_ = mapped;
     size_ = size;
@@ -363,9 +459,8 @@ Status MmapSharedSegmentBackend::CreateOwner(
 }
 
 Status MmapSharedSegmentBackend::ReserveLocal(
-    uint64_t size, const SharedSegmentOptions& /*options*/,
-    uintptr_t& base_addr) {
-    void* reserved = ReserveAlignedWindow(size, HugePageSize());
+    uint64_t size, const SharedSegmentOptions& options, uintptr_t& base_addr) {
+    void* reserved = ReserveAlignedWindow(size, MappingAlign(options));
     if (reserved == MAP_FAILED) {
         return Status::Memory(std::string("mmap reserve failed: ") +
                               std::strerror(errno));
@@ -404,13 +499,15 @@ Status MmapSharedSegmentBackend::ImportAndMap(
     }
 
     void* mapped = addr_;
-    auto status = MapMemFd(fd, size, mapped, mapped);
+    auto status = MapMemFd(fd, size, mapped, mapped, /*populate=*/false);
     close(fd);
     if (!status.ok()) {
         return status;
     }
     addr_ = mapped;
-    AdviseTransparentHugePages(addr_, size);
+    if (!options.hugetlb) {
+        AdviseTransparentHugePages(addr_, size);
+    }
     AdviseDontFork(addr_, size);
     status = RegisterMapped(options);
     if (!status.ok()) {