[he-yufeng]

Description

With the CUDA build, the TCP transport staged device memory through a synchronous cudaMemcpy on the legacy default stream, so any transfer of GPU buffers waited behind every kernel the process already had queued — the reporter sees streamed tokens reaching the client two at a time. The USE_MACA build already went through copyTcpCudaMemory (a private non-blocking stream, cudaMemcpyAsync, then a sync of that stream); the CUDA build didn't.

This lifts copyTcpCudaMemory out of the USE_MACA guard so every GPU build has it, and routes all four staging sites (ServerSession/ClientSession x readBody/writeBody) through it. The trailing stream sync keeps the copy's completion contract identical for callers.

Neighbor note: #4233 pipelines CUDA TCP staging with pinned buffers but deliberately keeps cudaStreamDefault (legacy-stream semantics) for its session stream, so the stall this fixes survives that PR. The two are complementary; whichever lands second rebases.

Fixes #4323

Module

• Transfer Engine (mooncake-transfer-engine)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:

# docker (mc-dev:local), mooncake-transfer-engine configured with -DBUILD_UNIT_TESTS=ON
ninja -C build tcp_transport_test tcp_write_visibility_test tcp_address_validation_test
./tcp_transport_test && ./tcp_write_visibility_test && ./tcp_address_validation_test
# CUDA-path compile check (no CUDA toolkit in this container): syntax-only
# compile of tcp_transport.cpp under -DUSE_CUDA against a minimal CUDA stub
g++ -fsyntax-only -std=c++20 -DUSE_CUDA -DUSE_TCP .../tcp_transport.cpp

Test results:

• Unit tests pass (52 + 52 + 10 on this head)
• Integration tests pass (if applicable)
• Manual testing done (CUDA-path syntax-only compile against a stub; no CUDA host here, so the runtime staging behavior on real GPU hardware is covered by the CI CUDA lanes, not by me)

Checklist

• I have performed a self-review of my own code
• I have run pre-commit on the files changed in this PR and all hooks pass
• I have added tests to prove my changes are effective (existing TCP staging/visibility suites cover the touched file; no new test because the behavior change is stream choice, not a new contract)
• My change is <500 LOC, no RFC needed

AI Assistance Disclosure

• AI tools were used (specify below)

Prepared with AI assistance (Kimi); the submitter reviewed the diff and the verification results above.

[he-yufeng]

Heads-up for reviewers: the reporter ran a wider repro matrix on this change over in #4323 and found one regression case (8 B reads while the sender is busy on a non-blocking stream: ~0.35 -> ~1.75 ms, pageable async D2H waiting behind an in-flight kernel). Analysis and the plan are in the issue thread: the trade still favors this fix (bounded ~1.4 ms vs an unbounded stall), and the residual is exactly what #4233's pinned staging closes, so whichever of the two lands second owns the re-measure.

[he-yufeng]

Merged current main (b6d5ba8) so the CUDA legs rerun on the fixed pipeline (#4513 restored the build, #4440 pinned the wheel-test torch). Merge was clean.

Rebuilt in the local docker toolchain: the staged-copy change compiles into the TCP transport, and tcp_transport_test passes against a live metadata server. The CUDA-gated staging test itself runs in the CI CUDA legs, which were previously red for the repo-wide Environ break, not this change.

[he-yufeng]

Open since 09-25 with a green matrix and no review yet, so a look from a transfer-engine owner when time allows: @alogfans @doujiang24 @chestnut-Q @staryxchen. The TCP transport staged CUDA buffers through a synchronous cudaMemcpy on the legacy default stream, so GPU transfers queued behind every kernel the process already had in flight; this routes all four staging sites through copyTcpCudaMemory (private stream, cudaMemcpyAsync) like the MACA build already did, with the same completion contract for callers.

[alogfans]

Sorry for delayed review :( However, The MLU backend does not provide the APIs required by the helper introduced in this PR.

copyTcpCudaMemory() was expanded from USE_MACA to all CUDA-like backends, but it calls cudaStreamCreateWithFlags(..., cudaStreamNonBlocking). The MLU shim only provides cudaStreamCreate, cudaMemcpyAsync, and cudaStreamSynchronize; it does not map cudaStreamCreateWithFlags or cudaStreamNonBlocking. As a result, builds with USE_MLU enabled will fail to compile.

[he-yufeng]

Thanks for the catch — you're right, the MLU shim only provides cnrtQueueCreate/cnrtMemcpyAsync_V3/cnrtQueueSync through the aliases in gpu_vendor/mlu.h, and neither cudaStreamCreateWithFlags nor cudaStreamNonBlocking is mapped.

Fixed in e30f420: the MLU path now uses plain cudaStreamCreate (which the shim maps to cnrtQueueCreate); the other backends keep the non-blocking-flagged stream, so the original "stay off the legacy stream" behavior is unchanged there. The blocking copy still ends in cudaStreamSynchronize, so the caller-visible completion contract is the same on both.

Verification, stated plainly: I don't have MLU headers here, so I can't run a real USE_MLU build. What I did run: the patched helper compiles under -Wall -Wextra -Werror against a stub exposing only the shim's actual API surface (create/memcpyAsync/sync/destroy), and the pre-fix form fails against that same stub with exactly the two missing names you listed. CI's CUDA leg covers the non-MLU branch, which is unchanged apart from the #ifdef placement.

[alogfans]

cudaStreamSynchronize(stream) waits for the copy, but does not ensure that the producer has finished writing src on another stream. Switching to cudaStreamNonBlocking may therefore expose a data race.

Please preserve producer-to-copy ordering, following [TENT NVLink’s synchronization](https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp#L98) (#3570). For TCP, establish this dependency at submission; synchronizing cudaStreamPerThread in the I/O thread would target the wrong stream.

[he-yufeng]

Addressed in 68dcab2faa. All three submission paths now establish the producer dependency on the submitting thread before any network I/O: one event per device and submission is recorded on cudaStreamPerThread and synchronized, rather than synchronizing the I/O thread's stream. CPU sources skip this path.

MLU now stays on its original synchronous cudaMemcpy, since that shim has no event or per-thread-stream API. The other backends keep the private nonblocking staging stream, with the submission-time dependency in front of it.

Added SubmissionWaitsForCallerProducerBeforeStaging: a gated producer runs on the same caller thread as submission, the test checks that submission cannot finish before the producer is released, and the received bytes must match the producer's final pattern. The existing legacy-stream nonblocking test remains as a separate control.

Validation: the three-file pre-commit run passed, including clang-format and codespell. I did not run a local build or CUDA test because available disk dropped below the 15 GiB local build gate, and I did not run a real USE_MLU build. CI on this head is the compile/CUDA verification gate; I am not claiming those environments as locally verified.

[he-yufeng]

Follow-up to the verification note: build-flags caught a CPU-only configuration mistake in 68dcab2faa. The sync helper was correctly hidden by the CUDA-like feature guards, but its three call sites were not, so the non-GPU build failed with TcpCallerStreamSync was not declared.

0e318551e05e puts all three call sites under exactly the same feature condition as the helper definition. The non-CUDA path is unchanged, and the touched-file pre-commit run passes. Fresh CI is running on this head.

[0z5a]

Hardware validation for this PR, since CI has no CUDA host: on real GPUs the legacy-stream staging does more than add latency. Through transfer_{write,read}_on_cuda it deadlocks, and this PR fixes it.

Why it deadlocks on the base: transfer_on_cuda_callback runs inside cudaLaunchHostFunc on the caller's stream and busy-waits for the batch. The TCP transport's synchronous cudaMemcpy (legacy default stream) then waits for that host function, which is waiting for the copy. With a legacy-default or blocking caller stream it never returns; a cudaStreamNonBlocking stream is unaffected. After this PR the staging copy runs on a private non-blocking stream, so the cycle is gone.

Setup: GH200 (aarch64, CUDA 13.2, driver 595), two nodes with one GPU each; MC_FORCE_TCP=1; 4 MiB GPU→GPU transfer_write_on_cuda; target verifies the bytes; 30 s watchdog. Builds from source: this PR at 20cbb068 vs its merge-base b140c1a9, plus current main (7d0ff6c6) for reference.

caller stream	path	merge-base b140c1a9	PR 20cbb068	main 7d0ff6c6
legacy default (stream_ptr=0)	same process (loopback segment)	hang	OK	hang
legacy default	two processes, same node	hang	OK (35 ms)	hang
legacy default	across nodes	hang	OK (68 ms)	—
cudaStreamNonBlocking	all three	OK	OK	OK

The same hang reproduces on 2× L40S (x86, CUDA 13.3) with the current nightly wheel, for both write and read, at 4 KiB and 64 MiB, with legacy and blocking streams (8/8). Minimal repro (no torch):

# Minimal repro: transfer_write_on_cuda over the TCP transport never completes
# when the caller passes the legacy default stream (or a blocking stream).
#   python repro_on_cuda_tcp.py legacy       -> hangs (watchdog prints HANG)
#   python repro_on_cuda_tcp.py nonblocking  -> OK
import ctypes, os, sys, threading

os.environ["MC_FORCE_TCP"] = "1"
from mooncake.engine import TransferEngine

kind = sys.argv[1] if len(sys.argv) > 1 else "legacy"
N = 4 << 20
rt = ctypes.CDLL("libcudart.so.13")
assert rt.cudaSetDevice(0) == 0
src, dst, s = ctypes.c_void_p(), ctypes.c_void_p(), ctypes.c_void_p()
assert rt.cudaMalloc(ctypes.byref(src), ctypes.c_size_t(N)) == 0
assert rt.cudaMalloc(ctypes.byref(dst), ctypes.c_size_t(N)) == 0
assert rt.cudaMemset(src, 0x5A, ctypes.c_size_t(N)) == 0
assert rt.cudaMemset(dst, 0, ctypes.c_size_t(N)) == 0
assert rt.cudaDeviceSynchronize() == 0
if kind != "legacy":  # cudaStreamNonBlocking = 1, default (blocking) = 0
    assert rt.cudaStreamCreateWithFlags(ctypes.byref(s), 1 if kind == "nonblocking" else 0) == 0

eng = TransferEngine()
assert eng.initialize("127.0.0.1", "P2PHANDSHAKE", "tcp", "") == 0
seg = f"127.0.0.1:{eng.get_rpc_port()}"
assert eng.register_memory(src.value, N, "cuda:0") == 0
assert eng.register_memory(dst.value, N, "cuda:0") == 0

threading.Timer(30, lambda: (print(f"{kind}: HANG after 30 s", flush=True), os._exit(3))).start()
eng.transfer_write_on_cuda(seg, src.value, dst.value, N, stream_ptr=s.value or 0)
assert rt.cudaStreamSynchronize(s) == 0
out = (ctypes.c_ubyte * N)()
assert rt.cudaMemcpy(out, dst, ctypes.c_size_t(N), 2) == 0
print(f"{kind}: OK" if bytes(out[:16]) == b"\x5a" * 16 and out[N - 1] == 0x5A else f"{kind}: CORRUPT", flush=True)
os._exit(0)

One gap worth closing here or in a follow-up: mooncake-wheel/tests/test_transfer_on_cuda.py hard-codes "rdma", "mlx5_bond_0" and skips otherwise, so this path never runs without an RDMA NIC. Running it with tcp and a parametrized caller stream (legacy / blocking / non-blocking) would have caught this.

AI assistance: the investigation and scripts were prepared with Claude; I reviewed them and ran the matrix on the hardware above.

[he-yufeng]

@alogfans The current head keeps MLU on the previous synchronous copy because its shim has neither stream flags nor the event/per-thread API. The other GPU paths synchronize cudaStreamPerThread once per source device at submission, before network I/O starts. That covers producer work on the caller's per-thread stream; it is not a claim that TransferRequest can describe an arbitrary producer stream.

transfer_*_on_cuda has a separate ordering guarantee: transfer_on_cuda_callback only runs after all preceding work on the stream passed to cudaLaunchHostFunc has completed. A producer on that same custom stream is therefore ordered before submission and before the private staging copy, even though TransferRequest itself has no stream field. A producer on an unrelated stream still needs caller-side synchronization.

The hardware evidence is narrower than a full producer-ordering matrix. On GH200, this PR completes the legacy/default cases that hang on its merge-base and current main, including the two-node case, while non-blocking callers remain OK. The separate 2x L40S result reproduces legacy and blocking hangs on the nightly/base build; it is not a validation run of this PR. None of these runs queued producer work on a custom stream, so I am not presenting them as proof of that ordering case. Could you re-review the implementation against these two contracts?