From b01f0e40c3edef446daa671ddf58ffc9f158546a Mon Sep 17 00:00:00 2001 From: "Wu, Yutong" Date: Sun, 18 Oct 2026 11:16:46 +0000 Subject: [PATCH] fix(umbp): pool the gather streams so the standalone server stops leaking them HbmCopyEngine created one HIP stream per (calling thread, device) for the gather kernel and never destroyed it, and LaunchDeviceGather kept its descriptor staging per thread the same way. The standalone server calls Submit from gRPC's sync server pool, which retires and spawns worker threads under load (at Kimi-K3 TP8 c14 the median worker lived 2.5 s and did 10 copies), so the streams grew without bound: 2370 after 37 minutes. Every HIP stream holds ~6.4 MiB of device memory for its kernarg pool, and with sglang at mem-fraction 0.9 the GPU filled up until hipStreamCreateWithFlags failed. HIP 7.2.4 then segfaults in its own failure path (HostQueue::terminate() releases a null virtual device), which is the "segfault at 8 ... in libamdhip64" that killed the server. The streams are now pooled per device and lent to one gather at a time, so concurrent copies still run on separate streams but a device's pool only grows to its peak number of concurrent gathers; the descriptor staging is kept per stream, which the caller holds exclusively until it synchronizes. A stream whose synchronize failed is dropped, as before. On the same workload (MI355X, ROCm 7.2.4) the unpatched server crashed after 15 minutes with 661 streams; with this change it served the full 3600 s window with 56 streams (7 per GPU) after 655 worker threads had come and gone. HbmCopyEngine::GatherStreamsCreated() lets a test assert the reuse; the new GatherStreamsOutliveTheirThreads test submits gathers from 16 successive threads and expects at most one new stream. --- .../distributed/transfer/hbm_copy_engine.cpp | 75 +++++++++++++------ src/umbp/include/umbp/common/device_gather.h | 3 + .../distributed/transfer/hbm_copy_engine.h | 6 ++ src/umbp/local/tiers/device_gather.hip | 33 ++++---- .../umbp/distributed/test_transfer_engine.cpp | 56 ++++++++++++++ 5 files changed, 139 insertions(+), 34 deletions(-) diff --git a/src/umbp/distributed/transfer/hbm_copy_engine.cpp b/src/umbp/distributed/transfer/hbm_copy_engine.cpp index a72cc204b..baec7e040 100644 --- a/src/umbp/distributed/transfer/hbm_copy_engine.cpp +++ b/src/umbp/distributed/transfer/hbm_copy_engine.cpp @@ -264,26 +264,55 @@ void LogPlanGeometry(const char* path, size_t plan_index, const TransferPlan& pl } } -// One non-blocking stream per device, per thread. Submit runs on PoolClient's -// executor threads, so a shared stream would serialize copies that the -// hipMemcpy path ran independently. -std::map& GatherStreams() { - static thread_local std::map* streams = new std::map(); - return *streams; -} +// Non-blocking gather streams, pooled per device and lent to one gather at a +// time, so concurrent copies still run on separate streams. They are not tied +// to the calling thread: the standalone server calls Submit from gRPC's sync +// pool, which retires and spawns threads under load, and every HIP stream holds +// ~6.4 MiB of device memory for its kernarg pool. Kept per thread, they leaked +// until the GPU was full and stream creation failed (which HIP 7.2 turns into a +// segfault). A device's pool grows only to its peak number of concurrent +// gathers. Never destroyed: worker threads can still be copying during static +// teardown. +class GatherStreamPool { + public: + static GatherStreamPool& Instance() { + static GatherStreamPool* pool = new GatherStreamPool(); + return *pool; + } -hipStream_t GatherStream(int device_id) { - auto& streams = GatherStreams(); - auto it = streams.find(device_id); - if (it != streams.end()) return it->second; - hipStream_t stream = nullptr; - if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) { - (void)hipGetLastError(); - return nullptr; + // The current device must already be `device_id`. nullptr if no stream is + // free and a new one cannot be created. + hipStream_t Acquire(int device_id) { + { + std::lock_guard lock(mutex_); + auto& idle = idle_[device_id]; + if (!idle.empty()) { + hipStream_t stream = idle.back(); + idle.pop_back(); + return stream; + } + } + hipStream_t stream = nullptr; + if (hipStreamCreateWithFlags(&stream, hipStreamNonBlocking) != hipSuccess) { + (void)hipGetLastError(); + return nullptr; + } + created_.fetch_add(1, std::memory_order_relaxed); + return stream; } - streams.emplace(device_id, stream); - return stream; -} + + void Release(int device_id, hipStream_t stream) { + std::lock_guard lock(mutex_); + idle_[device_id].push_back(stream); + } + + uint64_t Created() const { return created_.load(std::memory_order_relaxed); } + + private: + std::mutex mutex_; + std::map> idle_; + std::atomic created_{0}; +}; // Latches on the first gather failure and never tries again. A kernel that // cannot launch here will not start working later, and the fallback is always @@ -398,6 +427,8 @@ void* HbmCopyEngine::HostRegionDeviceAddress(const void* ptr, size_t size, int d return nullptr; } +uint64_t HbmCopyEngine::GatherStreamsCreated() { return GatherStreamPool::Instance().Created(); } + // Run the gather kernel over every eligible segment in the batch, and report // which plans it completed. // @@ -558,7 +589,7 @@ std::vector HbmCopyEngine::GatherEligiblePlans(const std::vector HbmCopyEngine::GatherEligiblePlans(const std::vector Submit(std::vector plans) override; + // Gather streams created since process start, across all engines. They are + // pooled per device and outlive the thread that first used them, so this + // tracks the peak number of gathers in flight at once, not the number of + // threads that have called Submit. Lets a test assert the reuse. + static uint64_t GatherStreamsCreated(); + private: // Why a plan did not take the gather kernel. The hot path never computes // this: it is filled in only when debug mode is on (UMBP_HBM_COPY_DEBUG), so diff --git a/src/umbp/local/tiers/device_gather.hip b/src/umbp/local/tiers/device_gather.hip index 5ba665d43..9372511f5 100644 --- a/src/umbp/local/tiers/device_gather.hip +++ b/src/umbp/local/tiers/device_gather.hip @@ -26,8 +26,9 @@ #include #include #include -#include +#include #include +#include #include "mori/utils/mori_log.hpp" #include "umbp/common/device_gather.h" @@ -71,25 +72,31 @@ __global__ void GatherFragmentsKernel(const DeviceGatherFragment* __restrict__ f for (size_t i = index; i < fragment.bytes; i += stride) dst[i] = src[i]; } -// Per-thread, per-device descriptor staging. +// Per-stream descriptor staging. // -// Per-thread because two threads serving the same device share one stream (see -// DeviceStream), so a shared staging buffer would be a data race even though -// the stream itself is ordered. Never freed, for the same reason DeviceStream -// never destroys its streams: releasing HIP resources from a thread or static -// destructor races with runtime teardown. The count is bounded by the server's -// worker threads times the devices they touch. +// The caller holds `stream` exclusively from the upload to the synchronize +// (see LaunchDeviceGather), so a stream's buffer is never in use twice at once, +// and the count is bounded by the caller's streams rather than by however many +// threads have ever called in. Never freed, for the same reason the streams are +// never destroyed: releasing HIP resources from a static destructor races with +// runtime teardown. struct DescriptorStaging { DeviceGatherFragment* host = nullptr; DeviceGatherFragment* device = nullptr; size_t capacity = 0; }; -DescriptorStaging* StagingFor(int device_id, size_t count) { - static thread_local std::map* staging_by_device = - new std::map(); +DescriptorStaging* StagingFor(hipStream_t stream, int device_id, size_t count) { + static std::mutex* staging_mutex = new std::mutex(); + static auto* staging_by_stream = new std::unordered_map(); - DescriptorStaging& staging = (*staging_by_device)[device_id]; + DescriptorStaging* staging_entry = nullptr; + { + // References into an unordered_map stay valid across rehashing. + std::lock_guard lock(*staging_mutex); + staging_entry = &(*staging_by_stream)[stream]; + } + DescriptorStaging& staging = *staging_entry; if (staging.capacity >= count) return &staging; // Grow with headroom so a slowly increasing batch size does not reallocate @@ -154,7 +161,7 @@ bool LaunchDeviceGather(const DeviceGatherFragment* fragments, size_t count, int (chunks + kThreadsPerBlock - 1) / kThreadsPerBlock, 1, kMaxBlocksPerFragment)); if (count > std::numeric_limits::max() / blocks_per_fragment) return false; - DescriptorStaging* staging = StagingFor(device_id, count); + DescriptorStaging* staging = StagingFor(stream, device_id, count); if (staging == nullptr) return false; std::memcpy(staging->host, fragments, count * sizeof(DeviceGatherFragment)); diff --git a/tests/cpp/umbp/distributed/test_transfer_engine.cpp b/tests/cpp/umbp/distributed/test_transfer_engine.cpp index 54e42aa42..592a520ff 100644 --- a/tests/cpp/umbp/distributed/test_transfer_engine.cpp +++ b/tests/cpp/umbp/distributed/test_transfer_engine.cpp @@ -26,6 +26,7 @@ #include #include +#include #include #include @@ -34,9 +35,11 @@ #include #include #include +#include #include #include "umbp/common/aligned_buffer.h" +#include "umbp/common/device_gather.h" #include "umbp/common/host_registration.h" #include "umbp/distributed/transfer/composite_transfer_engine.h" #include "umbp/distributed/transfer/hbm_copy_engine.h" @@ -292,6 +295,59 @@ TEST(HbmCopyEngine, BuilderSplitsAtHostRegistrationPieces) { ExpectPlanContract(items, planned); } +// Gathers submitted from a succession of short-lived threads, the way the +// standalone server's gRPC pool retires and spawns threads under load, share a +// pooled stream. A stream per thread grew until HIP could not create another. +TEST(HbmCopyEngine, GatherStreamsOutliveTheirThreads) { + int devices = 0; + if (hipGetDeviceCount(&devices) != hipSuccess || devices == 0) GTEST_SKIP() << "no GPU"; + if (!DeviceGatherEnabled()) GTEST_SKIP() << "gather kernel disabled"; + + constexpr size_t kBytes = 1 << 20; + constexpr size_t kFragment = 4096; + constexpr int kThreads = 16; + AlignedBuffer host(kBytes); + for (size_t i = 0; i < kBytes; ++i) host.data()[i] = static_cast(i * 131 + 7); + HbmCopyEngine engine; + engine.AddHostGatherRegion(host.data(), host.size()); + void* device = nullptr; + ASSERT_EQ(hipMalloc(&device, kBytes), hipSuccess); + ASSERT_EQ(hipMemset(device, 0, kBytes), hipSuccess); + const auto src = TransferRef::HostBytes(host.data(), host.size()); + const auto dst = TransferRef::HostBytes(device, kBytes, mori::io::MemoryLocationType::GPU, 0); + + const uint64_t launches_before = DeviceGatherLaunchCount(); + const uint64_t streams_before = HbmCopyEngine::GatherStreamsCreated(); + for (int t = 0; t < kThreads; ++t) { + std::thread([&, t] { + // Two fragments far apart, so they stay two segments: one is hipMemcpy's case. + const uint64_t first = t * 2 * kFragment; + const uint64_t second = kBytes / 2 + first; + auto planned = engine.Plan({MakeItem(src, first, dst, first, kFragment, 1), + MakeItem(src, second, dst, second, kFragment, 2)}); + std::vector failures; + if (auto handle = engine.Submit(std::move(planned.plans))) handle->Wait(&failures); + EXPECT_TRUE(failures.empty()); + }).join(); + } + if (DeviceGatherLaunchCount() == launches_before) { + (void)hipFree(device); + GTEST_SKIP() << "host region not gatherable on this device"; + } + EXPECT_EQ(DeviceGatherLaunchCount(), launches_before + kThreads); + EXPECT_LE(HbmCopyEngine::GatherStreamsCreated() - streams_before, 1u); + + std::vector copied(kBytes); + ASSERT_EQ(hipMemcpy(copied.data(), device, kBytes, hipMemcpyDeviceToHost), hipSuccess); + for (int t = 0; t < kThreads; ++t) { + for (const size_t offset : {t * 2 * kFragment, kBytes / 2 + t * 2 * kFragment}) { + EXPECT_EQ(std::memcmp(copied.data() + offset, host.data() + offset, kFragment), 0) + << "thread " << t << " offset " << offset; + } + } + (void)hipFree(device); +} + // --------------------------------------------------------------------------- // MoriIoEngine (planning only: the engine is never Init()ed) // ---------------------------------------------------------------------------