Repository navigation
fix(umbp): pool the gather streams so the standalone server stops leaking them - #799
Merged
Merged
Conversation
…king 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Change
The standalone server's
HbmCopyEnginekept one non-blocking HIP stream per (thread, device) for the gather kernel and never destroyed it. The server callsSubmitfrom gRPC's sync thread pool, which retires and spawns threads under load (a TP8 hour saw ~1,300 threads come and go), and every HIP stream holds ~6.4 MiB of device memory for its kernarg pool. Next to sglang atmem_fraction 0.9, the streams used up the remaining VRAM untilhipStreamCreateWithFlagsfailed, and HIP 7.2 turns that failure into a segfault (HostQueue::terminate()callsThread::Release()on a null virtual device:segfault at 8 ... in libamdhip64.so.7.2.70204). That is why only TP8 at high concurrency crashed.Gather streams are now pooled per device and lent to one gather at a time, so a device holds as many streams as its peak number of concurrent gathers, and the descriptor staging is kept per stream instead of per thread.
HbmCopyEngine::GatherStreamsCreated()exposes the count for the new testHbmCopyEngine.GatherStreamsOutliveTheirThreads: 16 short-lived threads, one gather each, at most one stream created (main fails it with 16 vs 1).The null dereference on a failed stream creation is a clr bug and will be reported to ROCm separately.
Verification
test_transfer_engineandtest_hbm_backendpass.test_umbp_pool_client_rangesgives the same per-case results as main (its remote cases need RDMA, which the test container lacks).fix/umbp-sa-prewarm-queuesbuilds on this branch and fixes the stalls the server then showed on amdgpu 6.16.