[GH-ISSUE #4986] [Bug]: TENT nvlink copies are routed through GPU 0 when the first write to a peer comes from a thread on another device (7x slower) #1264

New issue

Open

opened 2026-10-29 21:24:37 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 21:24:37 +00:00

Owner

Originally created by @leoleoasd on GitHub (Oct 27, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4986

Originally assigned to: @TTThanos on GitHub.

Bug Report

With TENT (MC_USE_TENT=1), an intra-node write from a buffer on GPU 1 to a peer process's buffer on GPU 2 runs at ~109 GB/s instead of ~771 GB/s whenever the thread that issues the first write to that peer buffer has CUDA current device 0, which is the default for every new thread.

A self-contained script (mooncake + torch only) is attached: tent_nvlink_ipc_device.py. It starts two processes on one node. The target registers a 1 GB buffer on --dst. The writer registers a buffer on --src and writes it into the target 10 times from a new thread, using transfer_submit_write / transfer_check_status.

MC_USE_TENT=1 python tent_nvlink_ipc_device.py            # thread on cuda:0 at first write
MC_USE_TENT=1 python tent_nvlink_ipc_device.py --bind     # thread calls torch.cuda.set_device(src) first
MC_USE_TENT=1 python tent_nvlink_ipc_device.py --src 0    # control: source on GPU 0

Results on B300 (NVLink 5), median of 1 GB writes 2..10:

Source -> target	Thread's device at first write	Median	Bandwidth
GPU 1 -> GPU 2	0	9.83 ms	109 GB/s
GPU 1 -> GPU 2	1 (`--bind`)	1.39 ms	771 GB/s
GPU 0 -> GPU 2	0	1.39 ms	771 GB/s

"""Minimal repro: Mooncake TENT nvlink copies are routed through GPU 0 when the first write to a peer is issued from a
thread whose CUDA current device is 0 (the default for any new thread).

NVLinkTransport::relocateSharedMemoryAddress (tent/src/transport/nvlink/nvlink_transport.cpp) opens the peer's
CUDA IPC handle with cudaIpcOpenMemHandle on the *calling thread's current device* and caches the mapping. If that
thread is on GPU 0 while the source buffer is on GPU 1, every later copy GPU 1 -> GPU 2 runs in GPU 0's context
(two NVLink hops, one driver stream), even though TENT submits it on a GPU 1 stream.

Two processes on one node, no other dependency than mooncake + torch:
  target: allocates and registers a buffer on --dst GPU
  writer: allocates and registers a buffer on --src GPU, then writes it into the target from a NEW thread
          --bind: the new thread calls torch.cuda.set_device(src) before its first write (workaround)

    MC_USE_TENT=1 python tent_nvlink_ipc_device.py            # unbound thread: ~110 GB/s on B300
    MC_USE_TENT=1 python tent_nvlink_ipc_device.py --bind     # bound thread:   ~780 GB/s
Run with --src 0 to see that GPU 0 as source is unaffected.
"""

import argparse
import multiprocessing as mp
import os
import socket
import threading
import time

MB = 1 << 20


def make_engine(gpu, nic):
    import torch
    from mooncake.engine import TransferEngine

    torch.cuda.set_device(gpu)
    ip = socket.gethostbyname(socket.gethostname())
    engine = TransferEngine()
    assert engine.initialize(ip, "P2PHANDSHAKE", "rdma", nic) == 0
    return engine, f"{ip}:{engine.get_rpc_port()}"


def target(gpu, nic, nbytes, info, done):
    import torch

    engine, session = make_engine(gpu, nic)
    buf = torch.zeros(nbytes, dtype=torch.uint8, device=f"cuda:{gpu}")
    assert engine.register_memory(buf.data_ptr(), nbytes) == 0
    info.put((session, buf.data_ptr()))
    done.wait()
    torch.cuda.synchronize()
    print(f"target cuda:{gpu}: first/last byte {int(buf[0])}/{int(buf[-1])} (expect 7/7)")


def writer(gpu, nic, nbytes, n, bind, info, done):
    import torch

    engine, _ = make_engine(gpu, nic)
    buf = torch.full((nbytes,), 7, dtype=torch.uint8, device=f"cuda:{gpu}")
    torch.cuda.synchronize()
    assert engine.register_memory(buf.data_ptr(), nbytes) == 0
    session, remote = info.get()
    times = []

    def work():
        if bind:
            torch.cuda.set_device(gpu)
        print(f"writer cuda:{gpu}: new thread, current device {torch.cuda.current_device()} at first write")
        for _ in range(n):
            t = time.perf_counter()
            batch = engine.transfer_submit_write(session, buf.data_ptr(), remote, nbytes)
            assert batch
            while (rc := engine.transfer_check_status(batch)) == 0:
                pass
            assert rc > 0, rc
            times.append(time.perf_counter() - t)

    th = threading.Thread(target=work)
    th.start()
    th.join()
    done.set()
    steady = sorted(times[1:])  # the first write includes opening the IPC handle
    print(f"writer cuda:{gpu} {'bind' if bind else 'no bind'}: {n} x {nbytes // MB} MB, "
          f"median {steady[len(steady) // 2] * 1e3:.2f} ms = {nbytes / steady[len(steady) // 2] / 1e9:.0f} GB/s")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--src", type=int, default=1)
    p.add_argument("--dst", type=int, default=2)
    p.add_argument("--nic", default="", help="RDMA device list for initialize() (not used by nvlink)")
    p.add_argument("--mb", type=int, default=1024)
    p.add_argument("--n", type=int, default=10)
    p.add_argument("--bind", action="store_true", help="set the writing thread's device to --src first")
    a = p.parse_args()
    os.environ.setdefault("MC_USE_TENT", "1")
    ctx = mp.get_context("spawn")
    info, done = ctx.Queue(), ctx.Event()
    procs = [ctx.Process(target=target, args=(a.dst, a.nic, a.mb * MB, info, done)),
             ctx.Process(target=writer, args=(a.src, a.nic, a.mb * MB, a.n, a.bind, info, done))]
    for pr in procs:
        pr.start()
    for pr in procs:
        pr.join()

Before submitting...

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

Originally created by @leoleoasd on GitHub (Oct 27, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4986 Originally assigned to: @TTThanos on GitHub. ### Bug Report With TENT (MC_USE_TENT=1), an intra-node write from a buffer on GPU 1 to a peer process's buffer on GPU 2 runs at ~109 GB/s instead of ~771 GB/s whenever the thread that issues the first write to that peer buffer has CUDA current device 0, which is the default for every new thread. A self-contained script (mooncake + torch only) is attached: `tent_nvlink_ipc_device.py`. It starts two processes on one node. The target registers a 1 GB buffer on `--dst`. The writer registers a buffer on `--src` and writes it into the target 10 times from a new thread, using `transfer_submit_write` / `transfer_check_status`. ```bash MC_USE_TENT=1 python tent_nvlink_ipc_device.py # thread on cuda:0 at first write MC_USE_TENT=1 python tent_nvlink_ipc_device.py --bind # thread calls torch.cuda.set_device(src) first MC_USE_TENT=1 python tent_nvlink_ipc_device.py --src 0 # control: source on GPU 0 ``` Results on B300 (NVLink 5), median of 1 GB writes 2..10: | Source -> target | Thread's device at first write | Median | Bandwidth | |---|---|---|---| | GPU 1 -> GPU 2 | 0 | 9.83 ms | 109 GB/s | | GPU 1 -> GPU 2 | 1 (`--bind`) | 1.39 ms | 771 GB/s | | GPU 0 -> GPU 2 | 0 | 1.39 ms | 771 GB/s | ``` """Minimal repro: Mooncake TENT nvlink copies are routed through GPU 0 when the first write to a peer is issued from a thread whose CUDA current device is 0 (the default for any new thread). NVLinkTransport::relocateSharedMemoryAddress (tent/src/transport/nvlink/nvlink_transport.cpp) opens the peer's CUDA IPC handle with cudaIpcOpenMemHandle on the *calling thread's current device* and caches the mapping. If that thread is on GPU 0 while the source buffer is on GPU 1, every later copy GPU 1 -> GPU 2 runs in GPU 0's context (two NVLink hops, one driver stream), even though TENT submits it on a GPU 1 stream. Two processes on one node, no other dependency than mooncake + torch: target: allocates and registers a buffer on --dst GPU writer: allocates and registers a buffer on --src GPU, then writes it into the target from a NEW thread --bind: the new thread calls torch.cuda.set_device(src) before its first write (workaround) MC_USE_TENT=1 python tent_nvlink_ipc_device.py # unbound thread: ~110 GB/s on B300 MC_USE_TENT=1 python tent_nvlink_ipc_device.py --bind # bound thread: ~780 GB/s Run with --src 0 to see that GPU 0 as source is unaffected. """ import argparse import multiprocessing as mp import os import socket import threading import time MB = 1 << 20 def make_engine(gpu, nic): import torch from mooncake.engine import TransferEngine torch.cuda.set_device(gpu) ip = socket.gethostbyname(socket.gethostname()) engine = TransferEngine() assert engine.initialize(ip, "P2PHANDSHAKE", "rdma", nic) == 0 return engine, f"{ip}:{engine.get_rpc_port()}" def target(gpu, nic, nbytes, info, done): import torch engine, session = make_engine(gpu, nic) buf = torch.zeros(nbytes, dtype=torch.uint8, device=f"cuda:{gpu}") assert engine.register_memory(buf.data_ptr(), nbytes) == 0 info.put((session, buf.data_ptr())) done.wait() torch.cuda.synchronize() print(f"target cuda:{gpu}: first/last byte {int(buf[0])}/{int(buf[-1])} (expect 7/7)") def writer(gpu, nic, nbytes, n, bind, info, done): import torch engine, _ = make_engine(gpu, nic) buf = torch.full((nbytes,), 7, dtype=torch.uint8, device=f"cuda:{gpu}") torch.cuda.synchronize() assert engine.register_memory(buf.data_ptr(), nbytes) == 0 session, remote = info.get() times = [] def work(): if bind: torch.cuda.set_device(gpu) print(f"writer cuda:{gpu}: new thread, current device {torch.cuda.current_device()} at first write") for _ in range(n): t = time.perf_counter() batch = engine.transfer_submit_write(session, buf.data_ptr(), remote, nbytes) assert batch while (rc := engine.transfer_check_status(batch)) == 0: pass assert rc > 0, rc times.append(time.perf_counter() - t) th = threading.Thread(target=work) th.start() th.join() done.set() steady = sorted(times[1:]) # the first write includes opening the IPC handle print(f"writer cuda:{gpu} {'bind' if bind else 'no bind'}: {n} x {nbytes // MB} MB, " f"median {steady[len(steady) // 2] * 1e3:.2f} ms = {nbytes / steady[len(steady) // 2] / 1e9:.0f} GB/s") if __name__ == "__main__": p = argparse.ArgumentParser() p.add_argument("--src", type=int, default=1) p.add_argument("--dst", type=int, default=2) p.add_argument("--nic", default="", help="RDMA device list for initialize() (not used by nvlink)") p.add_argument("--mb", type=int, default=1024) p.add_argument("--n", type=int, default=10) p.add_argument("--bind", action="store_true", help="set the writing thread's device to --src first") a = p.parse_args() os.environ.setdefault("MC_USE_TENT", "1") ctx = mp.get_context("spawn") info, done = ctx.Queue(), ctx.Event() procs = [ctx.Process(target=target, args=(a.dst, a.nic, a.mb * MB, info, done)), ctx.Process(target=writer, args=(a.src, a.nic, a.mb * MB, a.n, a.bind, info, done))] for pr in procs: pr.start() for pr in procs: pr.join() ``` ### Before submitting... - [ ] Ensure you searched for relevant issues and read the [documentation]

bsakaguchi added the

high priority

bug

labels

2026-10-29 21:24:37 +00:00

bsakaguchi commented

2026-10-29 21:24:37 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 27, 2026):

Thanks for opening this issue, @leoleoasd!

Field	Value
Issue	#4986
GitHub user ID	`37735580`
Reporter	@leoleoasd

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Oct 27, 2026): Thanks for opening this issue, @leoleoasd! | Field | Value | |-------|-------| | **Issue** | #4986 | | **GitHub user ID** | `37735580` | | **Reporter** | @leoleoasd | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:24:38 +00:00

Author

Owner

@he-yufeng commented on GitHub (Oct 28, 2026):

Root cause is one line narrower than the report, and it is in the policy rather than in relocateSharedMemoryAddress itself.

DefaultNvlinkTransportPolicy::openIpcMemHandle (mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp) calls cudaIpcOpenMemHandle(address, handle, cudaIpcMemLazyEnablePeerAccess) with no device set. cudaIpcOpenMemHandle maps into the calling thread's current context, so the mapping is bound to whatever device that thread happens to be on, which for a fresh thread is 0. Every later copy then runs a pointer that lives in GPU 0's context, hence the two hops you measured.

The policy already knows how to pick the right device for a copy. selectStreamDevice in the same class does getDeviceForPointer(local_buffer) and falls back to the current device, so the stream is on the source device while the mapping is not. And the MUSA policy already records where the handle was opened, shm_entry.device_id = ipc_device, guarded by #ifdef USE_MUSA; the CUDA policy passes opened_device = -1 and explicitly says the recorded device is only needed by context-sensitive closes. So the concept is present and deliberately unconnected on the CUDA path.

Two consequences for whoever fixes this, beyond setting the device before the open:

The cache key has to carry the device. remap_entries_ is keyed by (target_id, entry.addr). Once a handle is opened on a device it is only cheap from that device, so the same remote entry reached from two local devices needs two mappings and the key must distinguish them, otherwise the fix just moves the bug from "first thread wins" to "first device wins". Your --bind run hides this because binding makes every open land on one device.

Closing has to match. cudaIpcCloseMemHandle is per mapping, so a device-aware cache needs a device-aware close, and the existing opened_device out-param is the natural place to carry it. This is where I want to flag a collision before anyone starts.

Three open PRs are already restructuring exactly this function and its mapping lifetime, so I have deliberately not started an implementation here: #2380 fixes a VMM resource leak in relocateSharedMemoryAddress, #4340 rewrites the relocation for suballocated buffers and adds a relocate path over remap_entries_, and #4488 is its follow-up on IPC owner-device ordinals for HIP. None of them touches the calling-thread-device question above, so this bug is genuinely open and none of those three fixes it, but a device-aware cache key and a device-aware close land squarely on the lines they are editing. Whoever takes this should either build on top of #4340 or wait for that series to land.

For the record on severity: the report's numbers are self-consistent and the shape is the right one. A 7x penalty on intra-node NVLink only when the first write to a peer is issued from a thread on another device is easy to miss in a benchmark that sets the device up front, which is why it survives this long.

@he-yufeng commented on GitHub (Oct 28, 2026): Root cause is one line narrower than the report, and it is in the policy rather than in `relocateSharedMemoryAddress` itself. `DefaultNvlinkTransportPolicy::openIpcMemHandle` (`mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp`) calls `cudaIpcOpenMemHandle(address, handle, cudaIpcMemLazyEnablePeerAccess)` with no device set. `cudaIpcOpenMemHandle` maps into the calling thread's current context, so the mapping is bound to whatever device that thread happens to be on, which for a fresh thread is 0. Every later copy then runs a pointer that lives in GPU 0's context, hence the two hops you measured. The policy already knows how to pick the right device for a copy. `selectStreamDevice` in the same class does `getDeviceForPointer(local_buffer)` and falls back to the current device, so the stream is on the source device while the mapping is not. And the MUSA policy already records where the handle was opened, `shm_entry.device_id = ipc_device`, guarded by `#ifdef USE_MUSA`; the CUDA policy passes `opened_device = -1` and explicitly says the recorded device is only needed by context-sensitive closes. So the concept is present and deliberately unconnected on the CUDA path. Two consequences for whoever fixes this, beyond setting the device before the open: The cache key has to carry the device. `remap_entries_` is keyed by `(target_id, entry.addr)`. Once a handle is opened on a device it is only cheap from that device, so the same remote entry reached from two local devices needs two mappings and the key must distinguish them, otherwise the fix just moves the bug from "first thread wins" to "first device wins". Your `--bind` run hides this because binding makes every open land on one device. Closing has to match. `cudaIpcCloseMemHandle` is per mapping, so a device-aware cache needs a device-aware close, and the existing `opened_device` out-param is the natural place to carry it. This is where I want to flag a collision before anyone starts. Three open PRs are already restructuring exactly this function and its mapping lifetime, so I have deliberately not started an implementation here: #2380 fixes a VMM resource leak in `relocateSharedMemoryAddress`, #4340 rewrites the relocation for suballocated buffers and adds a `relocate` path over `remap_entries_`, and #4488 is its follow-up on IPC owner-device ordinals for HIP. None of them touches the calling-thread-device question above, so this bug is genuinely open and none of those three fixes it, but a device-aware cache key and a device-aware close land squarely on the lines they are editing. Whoever takes this should either build on top of #4340 or wait for that series to land. For the record on severity: the report's numbers are self-consistent and the shape is the right one. A 7x penalty on intra-node NVLink only when the first write to a peer is issued from a thread on another device is easy to miss in a benchmark that sets the device up front, which is why it survives this long.

bsakaguchi referenced this issue

2026-10-29 21:32:55 +00:00

[PR #1264] [MERGED] [Doc] Add update for RBG + SGLang HiCache integration #2116