kvcache-ai / kvcache-ai/Mooncake

[RFC]: TCP Transport Optimization for Non-RDMA Fallback

Open
#3,446 1 comment 1 reaction 0 assignees View on GitHub
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Summary

Make the classic TE `tcp` transport a usable high-performance fallback on
clusters without RDMA, not just a functional last resort.

Today's bottlenecks are structural, not protocol-level: grouped transfers are
strictly serialized (one slice in flight), the transport runs on a single
`io_context` thread, and connection setup blocks that thread. This RFC fixes
those three plus socket tuning and connection-pool hardening.

**No wire-format change** — all work is scheduling, threading, and socket
config; a patched side interoperates with an unpatched peer (v1 and v2, both
directions). Split into 2 PRs. Does not overlap `rdma_twosided` (#3377), which
targets RDMA-equipped clusters.

## Motivation

`tcp` is the lowest-priority transport (`selectTransport`: `tcp=1 < rdma=2`),
so it only carries traffic where RDMA is absent — exactly where it must not be
a cliff. Four problems:

1. **RTT-bound, not bandwidth-bound.** `startTransferSequence` advances one
slice at a time; under v2 the next slice waits for the peer's ack. Effective
throughput is `slice_size / RTT` regardless of link speed — 128KB over 100μs
caps at ~1.3GB/s.
2. **Single-threaded.** One worker handles `accept`, every inbound `readBody`,
and every outbound transfer. One core tops out at ~5–8GB/s.
3. **Blocking connect (defect-grade).** `getConnection` calls `resolve()` /
`asio::connect()` synchronously on the io thread; one DNS/SYN timeout
(1s/3s) freezes the whole transport, `accept` included.
4. **No congestion tuning.** Only `TCP_NODELAY` is set, so pipelining depends
entirely on `SO_SNDBUF ≥ BDP`, which default autotuning misses on high-BDP
links (100Gbps × 2ms = 25MB).

Two cheap side-fixes: the 17B `SessionHeader` is a separate `async_write` (a
tiny packet per slice under `TCP_NODELAY`), and CUDA staging does a
`new/delete` per chunk in all four directions.

**Not a problem:** the CPU path is already zero-copy (`writeBody` hands
`addr + offset` straight to `async_write`). Only the CUDA branch stages through
host memory. So "zero-copy" here means dropping unnecessary chunking and
reusing staging — not removing a copy.

## Non-goals

- Replacing RDMA as the high-bandwidth path — this is fallback QoS.
- Changing the wire protocol (v1/v2 framing, status frame, `kOpcodeV2Flag`
untouched); `v3` batching is deferred, see Open questions.
- `SO_ZEROCOPY` — needs source pages valid until the error-queue completion,
conflicting with releasing the buffer at `COMPLETED`; also a loss <256KB.
- Migrating off asio (`io_uring` / `coro_io`) — bottleneck is concurrency and
buffer sizing, not syscall cost.
- GPU copy pipelining — orthogonal, see Future work.

## Architecture

Current: single thread, window = 1, blocking connect.
Proposed: bounded concurrent window over N streams on an `io_context` pool.

```
submit* ─ StreamGroup{slices, atomic cursor} window = min(pending, N)
├─ runner_0 ─ socket_0 ─ slice_0 → slice_4 → ...
├─ runner_1 ─ socket_1 ─ slice_1 → slice_5 → ...
└─ runner_N ─ socket_N ─ ...

io_context pool (M threads, lazy) each socket pinned to one ctx
acquireConnection: pool hit → sync; miss → async_resolve/connect
```

Two invariants keep it cheap:

- **Per-socket serialization preserved.** Each socket binds to one
`io_context`; sessions already use `socket_->get_executor()` (not a global
context), so `Client/ServerSession` need no locks and no logic change. The
only global-context reference — the grouped continuation post — moves to the
socket executor.
- **One request per socket preserved.** Concurrency comes from N sockets, not
pipelining one socket. The server sees the same "header then payload"
sequence, so `ServerSession` is untouched and old peers interoperate.

**Adaptive window** (no workload mode): `window = min(pending, N)`. Large
transfers have few slices → narrows to 1–2 streams (avoids send-buffer
contention/incast); small transfers have many → fills N to hide ack RTT.
Chunking varies only by memory type (CPU: none; CUDA: staged), not by size.
Default `N=4` (4 × 4MB ≈ 16MB in flight, spans 0.3MB–12.5MB BDP). Dispatch is a
shared atomic cursor (work-stealing) to avoid stranding on uneven slices.

## Roadmap (2 PRs)

Benchmarking goes first: there is currently **no benchmark**, so without a
baseline PR2's gains are unquantifiable and regressions invisible. PR1 is also
low-risk (tuning), while PR2 bundles the concurrency work under one review.

**PR1 — Baseline + socket/pool tuning.**
- `tcp_transport_bench` (configurable slice size/count/streams; throughput +
P50/P99; `tc netem` RTT).
- Socket options: `SO_SNDBUF`/`SO_RCVBUF` (`MC_TCP_SOCK_BUF`, 4MB,
`0`=autotune), `TCP_QUICKACK`, optional `TCP_CONGESTION`.
- Pool on by default; `MC_TCP_POOL_MAX_PER_PEER` cap; `cleanupIdleConnections`
off the locked path onto a timer.
- *Validation:* existing suite + pool-saturation case; baseline at 4KB/128KB/8MB.

**PR2 — Concurrency: multi-stream, threading, async connect.**
- N-runner window (shared atomic cursor), with `submitTransfer`/`Task` routed
through it to **bound** their unlimited fire-and-forget loop.
- Header + first chunk coalesced into one `async_write` (wire bytes identical);
CPU path stops chunking; CUDA staging reused per session.
- `io_context` pool (`MC_TCP_IO_THREADS`, lazy), each socket pinned to one ctx;
`getConnection`→async `acquireConnection` (resolve cache; fixes the blocking-
connect defect); `session->initiate()` on the socket executor.
- *Validation:* v1↔v2 mixed-version matrix (preserve the
`readWriteAck()`/`writeBody()` order verbatim); failing `resolve`/`connect`
no longer stalls others; disconnect / in-flight-destruction cases; throughput
vs PR1 — small slices scale ~N, no regression at 8MB.

## Alternatives considered

| Option | Why not now |
| --- | --- |
| `SO_ZEROCOPY` | Source-page lifetime conflicts with `COMPLETED`; loss <256KB. |
| `io_uring` backend | Bottleneck isn't syscall cost; transport rewrite. |
| `coro_io` rebuild | Same, larger blast radius; asio yields these wins as-is. |
| `v3` batch header | Real for tiny slices but needs v1/v2/v3 compat; PR2's window covers most of it — decide from data. |
| Pipeline one socket | Needs on-wire request IDs + stateful server; N sockets gets the window with zero protocol change. |
| Everything in one PR | Benchmark must land first as the baseline; folding it in leaves the optimization unmeasured. |

## Resolved decisions

- **Defaults:** connection pooling defaults **on** (hazard closed by the PR1
cap; reuse removes per-transfer reconnect). `MC_TCP_SLICE_SIZE` **stays
64KB** — post-PR2 it only governs CUDA staging, so unchanged is a CPU no-op
and CUDA no-regression; opt into 1MB when wanted. Only the pool flip is a
documented behavior change.
- **`MC_TCP_STREAMS_PER_PEER=4`** fixed; no BDP autoscaling this round.
- **Config via `getenv`**, aligning with existing TCP style and avoiding merge
churn against active `rdma_*` work in `config.h`.

## Open questions

- **Is `v3` batching worth it?** Deferred until PR2 data shows a workload the
multi-stream window can't cover; unscheduled otherwise.

## Future work

- Pinned staging pool + copy/send double buffering (today's CUDA path is
synchronous `cudaMemcpy`→`async_write`→`cudaMemcpy`, no overlap, ~½ device
bandwidth). Orthogonal; ~150 LOC.
- `v3` multi-descriptor batching, if PR2 justifies it.

## Ask

Review (a) TCP-as-fallback being worth this investment, (b) the no-wire-change
boundary (`v3` deferred, not proposed here), and (c) the 2-PR order putting
benchmarking before the concurrency work.

---

Related: #3377 (`rdma_twosided`, a separate transport targeting RDMA-equipped clusters).

Contributor guide

Open the contributing guide

Research direction

Start by reading the existing tcp transport paths around startTransferSequence, getConnection, submitTransfer/Task, and the readWriteAck()/writeBody() ordering. Review the existing suite and establish the proposed tcp_transport_bench baseline before assessing the pool, socket tuning, and concurrency changes. Done means the two-PR plan is validated without wire-format changes, mixed v1/v2 compatibility is preserved, failed connects do not stall other work, and throughput improves without an 8MB regression.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, networking, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.