[GH-ISSUE #5042] [Feature Request]: [RFC] TCP transport: pinned GPU staging pool to overlap staging copies with socket IO #1277

New issue

Open

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

bsakaguchi commented

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

Owner

Originally created by @Czzzzz-5 on GitHub (Oct 28, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/5042

Summary

The TCP transport serializes synchronous 64 KiB staging copies with socket I/O, capping effective GPU-buffer transfer bandwidth at ~1.5 GB/s. This RFC proposes a persistent pinned-memory staging pool (one non-blocking stream + event per slot) for the TCP transport, overlapping staging copies with socket I/O, plus engine-side batched gather/scatter APIs for scattered GPU segments (e.g. vLLM KV caches). Measured 1.6–2.1x speedup on 128k-token KV transfers (2× RTX 5090, TCP loopback), raising effective bandwidth to ~3.1 GB/s.

Motivation

The TCP transport stages GPU buffers through a host staging buffer one
64 KiB chunk at a time: a synchronous cudaMemcpy, then
asio::async_write / async_read, then the next chunk
(ServerSession::writeBody, ClientSession::readBody). Within one
lane, staging copies and socket I/O are strictly serialized.

Measured with per-chunk timing instrumentation on a loopback
KV-transfer benchmark (2x RTX 5090): ~14% of a lane's timeline is
spent in staging copies while the socket idles, capping effective
bandwidth at ~1.5 GB/s.

#3562 (reuse TCP CUDA staging buffers across chunks) removes per-chunk
allocation but keeps the copy→IO serialization; #3724 shards sessions
across io_context threads. Neither removes the serialization within
a single transfer. (Related: #5012 recently added an opt-in pinned
staging pool to the io_uring transport — this RFC proposes the same
idea for the TCP transport.)

Proposal

Add a pinned-memory staging pool
(transport/tcp_transport/gpu_staging_pool.h): N persistent pinned
slots, each with a non-blocking stream + event.

Send side: a whole request is staged with a single event-ordered
asynchronous D2H copy, then chunks are streamed to the socket
directly from the pinned slot.
Receive side: bytes land in the slot and are copied back to GPU
with a single H2D copy before COMPLETED is reported, preserving the
"applied at the destination" semantics of the v2 protocol.
Engine-side batched APIs on top of the pool:
submit_gather_write (multi-segment GPU gather into a slot + single
contiguous write, e.g. a vLLM KV cache spread over blocks),
receive-side scatter, both using cudaMemcpyBatchAsync, plus
lane-parallel split writes to consume multiple TCP lanes for one
transfer.
Fallback preserved: pool disabled / request larger than one slot
/ all slots busy → existing per-chunk path (one-time WARNING, no
silent behavior change).

Env knobs: MC_TCP_GPU_STAGING (0 disables),
MC_TCP_GPU_STAGING_SLOT_MB (default 16 MiB),
MC_TCP_GPU_STAGING_SLOTS (default 8).

Benchmark

Real-shape KV transfer bench (P/D processes over TCP loopback, KV
layouts matching vLLM v1; 16 token levels x 3 iterations + warmup;
receiver-side data-pattern verification; same-day A/B against main @
e58ad953):

model (KV shape)	128k payload	main	with pool	speedup
K3 (24L MLA, 27.6 KB/tok)	3.6 GB	2056 ms	1295 ms	1.59x
DSV3 (61L MLA, 70.3 KB/tok)	9.2 GB	5963 ms	2896 ms	2.06x
Qwen2.5-7B (28L GQA, 56 KB/tok)	7.5 GB	4919 ms	2445 ms	2.01x

Effective bandwidth: ~1.5 -> ~3.1 GB/s.

Downstream evidence (vLLM PD-disaggregated serving with a connector
using these APIs): 100k-token prompt (5.7 GB KV) transfers 4/4 where
the legacy per-descriptor path deadlocks; network phase 2.4-2.9 s ->
0.31-0.41 s.

Scope of the accompanying PR

~920 LOC, mooncake-transfer-engine/src/transport/tcp_transport/ +
mooncake-integration/transfer_engine/. AI-assisted porting/benchmark
analysis; the human submitter has reviewed the full diff.

Before submitting a new issue...

Make sure you already searched for relevant issues and read the documentation

Originally created by @Czzzzz-5 on GitHub (Oct 28, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/5042 ## Summary The TCP transport serializes synchronous 64 KiB staging copies with socket I/O, capping effective GPU-buffer transfer bandwidth at ~1.5 GB/s. This RFC proposes a persistent pinned-memory staging pool (one non-blocking stream + event per slot) for the TCP transport, overlapping staging copies with socket I/O, plus engine-side batched gather/scatter APIs for scattered GPU segments (e.g. vLLM KV caches). Measured 1.6–2.1x speedup on 128k-token KV transfers (2× RTX 5090, TCP loopback), raising effective bandwidth to ~3.1 GB/s. ## Motivation The TCP transport stages GPU buffers through a host staging buffer one 64 KiB chunk at a time: a synchronous `cudaMemcpy`, then `asio::async_write` / `async_read`, then the next chunk (`ServerSession::writeBody`, `ClientSession::readBody`). Within one lane, staging copies and socket I/O are strictly serialized. Measured with per-chunk timing instrumentation on a loopback KV-transfer benchmark (2x RTX 5090): ~14% of a lane's timeline is spent in staging copies while the socket idles, capping effective bandwidth at ~1.5 GB/s. #3562 (reuse TCP CUDA staging buffers across chunks) removes per-chunk allocation but keeps the copy→IO serialization; #3724 shards sessions across io_context threads. Neither removes the serialization *within* a single transfer. (Related: #5012 recently added an opt-in pinned staging pool to the io_uring transport — this RFC proposes the same idea for the TCP transport.) ## Proposal Add a pinned-memory staging pool (`transport/tcp_transport/gpu_staging_pool.h`): N persistent pinned slots, each with a non-blocking stream + event. - **Send side**: a whole request is staged with a single event-ordered asynchronous D2H copy, then chunks are streamed to the socket directly from the pinned slot. - **Receive side**: bytes land in the slot and are copied back to GPU with a single H2D copy before COMPLETED is reported, preserving the "applied at the destination" semantics of the v2 protocol. - **Engine-side batched APIs** on top of the pool: `submit_gather_write` (multi-segment GPU gather into a slot + single contiguous write, e.g. a vLLM KV cache spread over blocks), receive-side scatter, both using `cudaMemcpyBatchAsync`, plus lane-parallel split writes to consume multiple TCP lanes for one transfer. - **Fallback preserved**: pool disabled / request larger than one slot / all slots busy → existing per-chunk path (one-time WARNING, no silent behavior change). Env knobs: `MC_TCP_GPU_STAGING` (`0` disables), `MC_TCP_GPU_STAGING_SLOT_MB` (default 16 MiB), `MC_TCP_GPU_STAGING_SLOTS` (default 8). ## Benchmark Real-shape KV transfer bench (P/D processes over TCP loopback, KV layouts matching vLLM v1; 16 token levels x 3 iterations + warmup; receiver-side data-pattern verification; same-day A/B against main @ e58ad953): | model (KV shape) | 128k payload | main | with pool | speedup | |---|---|---|---|---| | K3 (24L MLA, 27.6 KB/tok) | 3.6 GB | 2056 ms | 1295 ms | 1.59x | | DSV3 (61L MLA, 70.3 KB/tok) | 9.2 GB | 5963 ms | 2896 ms | 2.06x | | Qwen2.5-7B (28L GQA, 56 KB/tok) | 7.5 GB | 4919 ms | 2445 ms | 2.01x | Effective bandwidth: ~1.5 -> ~3.1 GB/s. Downstream evidence (vLLM PD-disaggregated serving with a connector using these APIs): 100k-token prompt (5.7 GB KV) transfers 4/4 where the legacy per-descriptor path deadlocks; network phase 2.4-2.9 s -> 0.31-0.41 s. ## Scope of the accompanying PR ~920 LOC, `mooncake-transfer-engine/src/transport/tcp_transport/` + `mooncake-integration/transfer_engine/`. AI-assisted porting/benchmark analysis; the human submitter has reviewed the full diff. ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi commented

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

Author

Owner

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

Thanks for opening this issue, @Czzzzz-5!

Field	Value
Issue	#5042
GitHub user ID	`293024653`
Reporter	@Czzzzz-5

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 28, 2026): Thanks for opening this issue, @Czzzzz-5! | Field | Value | |-------|-------| | **Issue** | #5042 | | **GitHub user ID** | `293024653` | | **Reporter** | @Czzzzz-5 | 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:45 +00:00

Author

Owner

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

Two things from working on the same staging path, since they change the baseline this RFC is measured against and one of them is a correctness constraint rather than a performance one.

The synchronous cudaMemcpy described in the motivation is on its way out. #4324 moves the staging copy onto a private non-blocking stream and keeps a sync only for the copy's own completion contract, so it no longer queues behind every kernel on the caller's lane. That is a latency and correctness fix rather than the overlap this RFC is after: staging and socket I/O are still serialized, so the gap you measured is still there and the pool idea still composes on top of it. Worth reading the timing numbers as measured on the current synchronous path, since the split between copy time and socket idle time will shift once the copy is off the legacy stream. There is also a hard failure behind this: on real GPUs the legacy-stream staging deadlocks through transfer_{write,read}_on_cuda, which is what #4324 is fixing.

The constraint to carry into the pool design is producer ordering. A private stream does not inherit ordering from whatever wrote the source buffer, so a slot's copy has to establish that dependency before it reads src, or the pool turns a latency bug into a silent read of stale data. The pattern used in TENT NVLink is an event recorded and synchronized on the submitting thread's per-thread stream, once per device and submission, before any network I/O starts; SubmissionWaitsForCallerProducerBeforeStaging in tcp_cuda_staging_test.cpp holds a gate on that stream and asserts both that submission waits and that the staged bytes match. A per-slot stream plus event makes this easier to get right than the current single staging buffer, since each slot can own its dependency instead of relying on the lane's ordering.

The other half of the motivation, batched gather and scatter for scattered GPU segments, is separate enough from the staging pool that it might be worth splitting. The pool is confined to the TCP transport session; the gather and scatter APIs are engine-side and touch the request surface.

@he-yufeng commented on GitHub (Oct 28, 2026): Two things from working on the same staging path, since they change the baseline this RFC is measured against and one of them is a correctness constraint rather than a performance one. The synchronous `cudaMemcpy` described in the motivation is on its way out. #4324 moves the staging copy onto a private non-blocking stream and keeps a sync only for the copy's own completion contract, so it no longer queues behind every kernel on the caller's lane. That is a latency and correctness fix rather than the overlap this RFC is after: staging and socket I/O are still serialized, so the gap you measured is still there and the pool idea still composes on top of it. Worth reading the timing numbers as measured on the current synchronous path, since the split between copy time and socket idle time will shift once the copy is off the legacy stream. There is also a hard failure behind this: on real GPUs the legacy-stream staging deadlocks through `transfer_{write,read}_on_cuda`, which is what #4324 is fixing. The constraint to carry into the pool design is producer ordering. A private stream does not inherit ordering from whatever wrote the source buffer, so a slot's copy has to establish that dependency before it reads `src`, or the pool turns a latency bug into a silent read of stale data. The pattern used in TENT NVLink is an event recorded and synchronized on the submitting thread's per-thread stream, once per device and submission, before any network I/O starts; `SubmissionWaitsForCallerProducerBeforeStaging` in `tcp_cuda_staging_test.cpp` holds a gate on that stream and asserts both that submission waits and that the staged bytes match. A per-slot stream plus event makes this easier to get right than the current single staging buffer, since each slot can own its dependency instead of relying on the lane's ordering. The other half of the motivation, batched gather and scatter for scattered GPU segments, is separate enough from the staging pool that it might be worth splitting. The pool is confined to the TCP transport session; the gather and scatter APIs are engine-side and touch the request surface.

bsakaguchi referenced this issue

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

[PR #1277] [MERGED] [TENT] Add backward compatibility with TE #2126