kvcache-ai / kvcache-ai/Mooncake
[Bug]: HIP IPC transfers silently return the wrong bytes for sub-allocated device buffers
- 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
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