[GH-ISSUE #2016] [Bug]: HipTransport::setupP2PAccess leaves wrong active HIP device, causing OOM cascade in TP workers #560

New issue

Closed

opened 2026-10-29 21:16:26 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 21:16:26 +00:00

Owner

Originally created by @Lzy17 on GitHub (Apr 30, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2016

Originally assigned to: @amd-arozanov on GitHub.

Bug Report

Summary

When mooncake::HipTransport is constructed inside a multi-process tensor-parallel worker (one process per GPU), setupP2PAccess(int num_devices) walks every visible GPU with hipSetDevice(i) to enable peer access, but never restores the device that was active on entry. After the constructor returns, every TP rank's worker thread has its active HIP device pinned to num_devices - 1 (the highest-numbered GPU). All subsequent torch allocations from those ranks then land on that single GPU, exhausting its VRAM and crashing the process at model load.

Affected version

Repo: kvcache-ai/Mooncake
Branch: main
Commit: b0bda8caa4e729208d3d8c073091851ab5f55ac1
Build flag: -DUSE_HIP=ON
File: mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp
Function: static void setupP2PAccess(int num_devices)

Environment

8x AMD Instinct MI355X (gfx950) per node
ROCm 7.2 / libamdhip64.so.7
PyTorch with HIP backend
SGLang 0.5.10.post1 (image rocm/sgl-dev:v0.5.10.post1-rocm720-mi35x-20260428)
Tensor parallelism = 8 (one worker process per GPU)
HipTransport instantiated by SGLang's mooncake KV-disaggregation connector

Reproduce

Build upstream Mooncake with HIP enabled:

git clone https://github.com/kvcache-ai/Mooncake.git
cd Mooncake
bash dependencies.sh
mkdir build && cd build
cmake -DUSE_HIP=ON ..
make -j
make install     # writes engine.cpython-310-x86_64-linux-gnu.so + store.cpython-310-x86_64-linux-gnu.so
                 # into the active python site-packages/mooncake/

Launch any tensor-parallel server that constructs HipTransport per rank. Minimal repro using SGLang PD-disaggregation:

export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
python3 -m sglang.launch_server \
  --model-path <any-large-model> \
  --tensor-parallel-size 8 \
  --disaggregation-mode prefill \
  --disaggregation-transfer-backend mooncake \
  --disaggregation-ib-device <ib-devs> \
  --mem-fraction-static 0.6

Server fails to start. Each TP rank's log shows the diagnostic warning followed by the OOM:

WARNING: current device is not 2, but 7, which may cause useless memory allocation for torch CUDA context
...
torch.OutOfMemoryError: HIP out of memory. Tried to allocate 28.00 MiB.
GPU 7 has a total capacity of 287.98 GiB of which 0 bytes is free.

GPU 7 hits 100% VRAM use while GPUs 0–6 stay near zero, confirming all ranks are allocating onto the same device.

Root cause

setupP2PAccess (current upstream):

static void setupP2PAccess(int num_devices) {
    for (int i = 0; i < num_devices; ++i) {
        if (!checkHip(hipSetDevice(i), "...")) continue;
        for (int j = 0; j < num_devices; ++j) {
            if (i == j) continue;
            int can_access_peer = 0;
            hipDeviceCanAccessPeer(&can_access_peer, i, j);
            if (can_access_peer) hipDeviceEnablePeerAccess(j, 0);
        }
    }
    // <-- no restore: active device is now (num_devices - 1)
}

hipSetDevice is per-thread state. The HipTransport constructor is called from the worker process / thread that PyTorch later uses for its own allocator. Because the active device is never restored, every worker leaves this function pinned to GPU 7, and every torch.empty(...) / weight load from that worker targets GPU 7 instead of the rank's own GPU.

Suggested fix

Save the entry device with hipGetDevice and restore it before returning. This makes the function transparent to the caller's HIP context. (Bonus: clear sticky hipErrorPeerAccessAlreadyEnabled from the runtime's last-error slot so it doesn't surface on the next unrelated HIP call from PyTorch.)

static void setupP2PAccess(int num_devices) {
    int original_device = -1;
    if (!checkHip(hipGetDevice(&original_device),
                  "HipTransport: failed to get current device for restore")) {
        original_device = -1;
    }

    for (int i = 0; i < num_devices; ++i) {
        if (!checkHip(hipSetDevice(i), "HipTransport: failed to set device")) {
            continue;
        }
        for (int j = 0; j < num_devices; ++j) {
            if (i == j) continue;
            int can_access_peer = 0;
            if (!checkHip(hipDeviceCanAccessPeer(&can_access_peer, i, j),
                          "HipTransport: failed to query peer access")) {
                continue;
            }
            if (can_access_peer) {
                hipError_t result = hipDeviceEnablePeerAccess(j, 0);
                if (result == hipErrorPeerAccessAlreadyEnabled ||
                    result != hipSuccess) {
                    // swallow the sticky last-error so PyTorch's next HIP call
                    // does not observe a stale error
                    (void)hipGetLastError();
                }
            }
        }
    }

    // Restore the entry device so the calling thread is unchanged
    if (original_device >= 0) {
        (void)hipSetDevice(original_device);
    }
}

Verification

After applying the fix, rebuilding (make && make install), and relaunching the same TP=8 server:

All 8 ranks reach HTTP server is started.
VRAM is balanced across all 8 GPUs (~156 GB each on the prefill node, ~201 GB each on the decode node).
The current device is not X, but 7 warning no longer appears.
No OOM at startup.

Notes

The bug only manifests when HipTransport is constructed in a process that has its own meaningful active device (e.g. one TP worker per GPU). A single-process launcher that always re-pins its device after Mooncake init would not see it, which may be why CI did not catch this.
The CUDA twin (CudaTransport::setupP2PAccess, if present) should be checked for the same pattern.

Before submitting...

Ensure you searched for relevant issues and read the [documentation]

Originally created by @Lzy17 on GitHub (Apr 30, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/2016 Originally assigned to: @amd-arozanov on GitHub. ### Bug Report ## Summary When `mooncake::HipTransport` is constructed inside a multi-process tensor-parallel worker (one process per GPU), `setupP2PAccess(int num_devices)` walks every visible GPU with `hipSetDevice(i)` to enable peer access, but never restores the device that was active on entry. After the constructor returns, every TP rank's worker thread has its active HIP device pinned to `num_devices - 1` (the highest-numbered GPU). All subsequent `torch` allocations from those ranks then land on that single GPU, exhausting its VRAM and crashing the process at model load. ## Affected version - Repo: `kvcache-ai/Mooncake` - Branch: `main` - Commit: `b0bda8caa4e729208d3d8c073091851ab5f55ac1` - Build flag: `-DUSE_HIP=ON` - File: `mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp` - Function: `static void setupP2PAccess(int num_devices)` ## Environment - 8x AMD Instinct MI355X (`gfx950`) per node - ROCm 7.2 / `libamdhip64.so.7` - PyTorch with HIP backend - SGLang `0.5.10.post1` (image `rocm/sgl-dev:v0.5.10.post1-rocm720-mi35x-20260428`) - Tensor parallelism = 8 (one worker process per GPU) - HipTransport instantiated by SGLang's mooncake KV-disaggregation connector ## Reproduce 1. Build upstream Mooncake with HIP enabled: ```bash git clone https://github.com/kvcache-ai/Mooncake.git cd Mooncake bash dependencies.sh mkdir build && cd build cmake -DUSE_HIP=ON .. make -j make install # writes engine.cpython-310-x86_64-linux-gnu.so + store.cpython-310-x86_64-linux-gnu.so # into the active python site-packages/mooncake/ ``` 2. Launch any tensor-parallel server that constructs `HipTransport` per rank. Minimal repro using SGLang PD-disaggregation: ```bash export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 -m sglang.launch_server \ --model-path <any-large-model> \ --tensor-parallel-size 8 \ --disaggregation-mode prefill \ --disaggregation-transfer-backend mooncake \ --disaggregation-ib-device <ib-devs> \ --mem-fraction-static 0.6 ``` 3. Server fails to start. Each TP rank's log shows the diagnostic warning followed by the OOM: ``` WARNING: current device is not 2, but 7, which may cause useless memory allocation for torch CUDA context ... torch.OutOfMemoryError: HIP out of memory. Tried to allocate 28.00 MiB. GPU 7 has a total capacity of 287.98 GiB of which 0 bytes is free. ``` GPU 7 hits 100% VRAM use while GPUs 0–6 stay near zero, confirming all ranks are allocating onto the same device. ## Root cause `setupP2PAccess` (current upstream): ```cpp static void setupP2PAccess(int num_devices) { for (int i = 0; i < num_devices; ++i) { if (!checkHip(hipSetDevice(i), "...")) continue; for (int j = 0; j < num_devices; ++j) { if (i == j) continue; int can_access_peer = 0; hipDeviceCanAccessPeer(&can_access_peer, i, j); if (can_access_peer) hipDeviceEnablePeerAccess(j, 0); } } // <-- no restore: active device is now (num_devices - 1) } ``` `hipSetDevice` is per-thread state. The HipTransport constructor is called from the worker process / thread that PyTorch later uses for its own allocator. Because the active device is never restored, every worker leaves this function pinned to GPU 7, and every `torch.empty(...)` / weight load from that worker targets GPU 7 instead of the rank's own GPU. ## Suggested fix Save the entry device with `hipGetDevice` and restore it before returning. This makes the function transparent to the caller's HIP context. (Bonus: clear sticky `hipErrorPeerAccessAlreadyEnabled` from the runtime's last-error slot so it doesn't surface on the next unrelated HIP call from PyTorch.) ```cpp static void setupP2PAccess(int num_devices) { int original_device = -1; if (!checkHip(hipGetDevice(&original_device), "HipTransport: failed to get current device for restore")) { original_device = -1; } for (int i = 0; i < num_devices; ++i) { if (!checkHip(hipSetDevice(i), "HipTransport: failed to set device")) { continue; } for (int j = 0; j < num_devices; ++j) { if (i == j) continue; int can_access_peer = 0; if (!checkHip(hipDeviceCanAccessPeer(&can_access_peer, i, j), "HipTransport: failed to query peer access")) { continue; } if (can_access_peer) { hipError_t result = hipDeviceEnablePeerAccess(j, 0); if (result == hipErrorPeerAccessAlreadyEnabled || result != hipSuccess) { // swallow the sticky last-error so PyTorch's next HIP call // does not observe a stale error (void)hipGetLastError(); } } } } // Restore the entry device so the calling thread is unchanged if (original_device >= 0) { (void)hipSetDevice(original_device); } } ``` ## Verification After applying the fix, rebuilding (`make && make install`), and relaunching the same TP=8 server: - All 8 ranks reach `HTTP server is started`. - VRAM is balanced across all 8 GPUs (~156 GB each on the prefill node, ~201 GB each on the decode node). - The `current device is not X, but 7` warning no longer appears. - No OOM at startup. ## Notes - The bug only manifests when `HipTransport` is constructed in a process that has its own meaningful active device (e.g. one TP worker per GPU). A single-process launcher that always re-pins its device after Mooncake init would not see it, which may be why CI did not catch this. - The CUDA twin (`CudaTransport::setupP2PAccess`, if present) should be checked for the same pattern. ### Before submitting... - [x] Ensure you searched for relevant issues and read the [documentation]

bsakaguchi

2026-10-29 21:16:26 +00:00