kvcache-ai / kvcache-ai/Mooncake

[Bug]: HIP IPC transfers silently return the wrong bytes for sub-allocated device buffers

Open
#3,546 2 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Bug Report

## Environment

- 8× AMD Instinct MI308X (gfx942), ROCm 7.2.0
- PyTorch 2.9.1 (`torch.version.hip == 7.2.26015-fc0010cf6a`), Python 3.10.9
- Mooncake `main` @ `f90ae69`, built with `-DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx942`
- `protocol="hip"`, `P2PHANDSHAKE`, single node, intra-node GPU↔GPU

## What happens

When several device buffers that live inside **one** allocation are registered
individually, a peer reading them back gets the wrong data. Every
`transfer_sync_read` returns `0`, the byte counts match, nothing is logged — only
the contents are wrong. Each buffer comes back holding the *first* buffer's bytes.

We ran into this in a disaggregated RL weight-transfer path that registers each
FSDP2 parameter shard separately so inference workers can read straight into
their own parameter memory. It first looked like partial corruption
(e.g. `3141269/4194304 elements differ`), because the parameters that happened to
land at an allocation base were correct and the rest were not.

## Steps to reproduce

Two processes are required. A loopback transfer inside a single process never
opens an IPC handle, so this cannot be observed that way — which is also why
`mooncake-wheel/tests/test_transfer_on_hip.py` passes on an affected build.

```python
import multiprocessing as mp
import torch
from mooncake.engine import TransferEngine

N = 1 << 20 # bytes per buffer

def source(meta, done):
torch.cuda.set_device(0)
pool = torch.empty(4 * N, dtype=torch.uint8, device="cuda:0")
views = [pool[i * N : (i + 1) * N] for i in range(4)]
for i, v in enumerate(views):
v.fill_(i + 1) # buffer i holds the byte i+1
torch.cuda.synchronize()

e = TransferEngine()
assert e.initialize("127.0.0.1", "P2PHANDSHAKE", "hip", "") == 0
for v in views:
assert e.register_memory(v.data_ptr(), N) == 0
meta.put((f"127.0.0.1:{e.get_rpc_port()}", [v.data_ptr() for v in views]))
done.get(timeout=120)

def sink(meta, done):
torch.cuda.set_device(0)
segment, addresses = meta.get(timeout=120)

e = TransferEngine()
assert e.initialize("127.0.0.1", "P2PHANDSHAKE", "hip", "") == 0
pool = torch.zeros(4 * N, dtype=torch.uint8, device="cuda:0")
views = [pool[i * N : (i + 1) * N] for i in range(4)]
for v in views:
assert e.register_memory(v.data_ptr(), N) == 0
codes = [
e.transfer_sync_read(segment, v.data_ptr(), a, N)
for v, a in zip(views, addresses)
]
torch.cuda.synchronize()
print("return codes:", codes)
print("expected [1, 2, 3, 4], got", [int(v[0]) for v in views])
done.put(True)

if __name__ == "__main__":
ctx = mp.get_context("spawn")
meta, done = ctx.Queue(), ctx.Queue()
procs = [ctx.Process(target=f, args=(meta, done)) for f in (source, sink)]
for p in procs:
p.start()
for p in procs:
p.join()
```

Observed:

```
return codes: [0, 0, 0, 0]
expected [1, 2, 3, 4], got [1, 1, 1, 1]
```

Expected: `got [1, 2, 3, 4]`.

The four views above make the sub-allocation explicit, but nothing here depends
on that: four separately allocated small tensors reproduce it just as well,
because PyTorch's caching allocator places them in one segment anyway. Measured
with `hipMemGetAddressRange` on four separately allocated 2 MiB tensors:

```
buf0 offset=0 at allocation base
buf1 offset=2097152 NOT at base (all four inside one 20 MiB allocation)
buf2 offset=4194304 NOT at base
buf3 offset=6291456 NOT at base
```

`batch_register_memory` behaves the same, since `registerLocalMemoryBatch` loops
over `registerLocalMemory`.

## What we have narrowed it down to

`hipIpcGetMemHandle` always exports the whole allocation, but
`HipTransport::registerLocalMemory` records the caller's sub-range address as the
buffer base (`hip_transport.cpp:688,696`):

```cpp
hipIpcGetMemHandle(&handle, addr); // handle covers the whole allocation
desc.addr = (uint64_t)addr; // but we claim the buffer starts here
desc.length = length;
```

The reader relocates with (`hip_transport.cpp:772`):

```cpp
dest_addr = dest_addr - entry.addr + (uint64_t)shm_addr;
```

`shm_addr` maps the whole allocation while `entry.addr` is the sub-range, so every
sub-range collapses onto the allocation base. That matches the observed output
exactly: buffers 1–3 return buffer 0's bytes, and buffer 0 is correct because its
offset happens to be zero.

This looks like the same family as #1622, #1831 and #2035, where the CUDA paths
were changed to register the allocation base instead of the requested range. It
seems distinct from #2684 / #2752, which are about multi-protocol `rdma+hip`
transport selection and segment-descriptor failures rather than silently wrong
payloads.

### Before submitting...

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

Contributor guide

Open the contributing guide

Research direction

Start in hip_transport.cpp around lines 688, 696, and 772, then compare the related CUDA fixes from issues #1622, #1831, and #2035. Reproduce with the two-process example rather than the loopback test, and verify that registered sub-allocated buffers return [1, 2, 3, 4] while transfer return codes remain zero.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, pytorch
Domain
distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
64/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.