[GH-ISSUE #4897] [RFC]: Same-Rail KV Cache Transfer via a Local Proxy GPU for Classic Transfer Engine #1244

Open
opened 2026-10-29 21:24:22 +00:00 by bsakaguchi · 1 comment
Owner

Originally created by @yunjuanyunshu0902-zyh on GitHub (Oct 23, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4897

Changes proposed

Proposal

Introduce a default-off cross-rail avoidance path in the classic Transfer Engine (TE) RDMA transport. When a slice originates on a GPU whose same-rail NIC is not usable for the chosen remote peer NIC, the sending host stages the payload on a locally reachable proxy GPU that shares a PCIe switch (and NVLink) with the source GPU and whose same-rail NIC can reach the remote peer NIC on the same rail. The transport issues one cudaMemcpyAsync from the source GPU to a pinned proxy-GPU slot, then issues the RDMA WRITE from the proxy GPU through its same-rail NIC. Receiver semantics, upper-layer transfer completion, and endpoint lifecycles remain unchanged.

The goal is to remove cross-rail traffic from the spine fabric in PD-disaggregated inference. Spine switches are shared with other tenants, and cross-rail KV cache bursts observed in production trigger PFC pause-frame storms that propagate past the originating flow and degrade TTFT of co-located workloads. Keeping KV transfer traffic on the same rail eliminates the spine as a congestion domain for this workload without changing fabric configuration.

This proposal is scoped to the classic Transfer Engine RDMA transport. It does not change ECN or PFC settings, does not alter the TCP, EFA, UB, or TENT transports, and does not take ownership of peer liveness or cluster membership decisions.

Relationship to Existing Work

PXN relaying was first described in the NCCL communication library (NCCL PXN): for a cross-rail GPU pair, egress through a locally reachable proxy GPU whose same-rail NIC can reach the remote peer on the same rail. This proposal applies the same locality principle to the Mooncake classic Transfer Engine data path. It is implemented as an opt-in build feature, keeps the existing RDMA worker pool as the submission authority, and reuses the existing endpoint, metadata, and completion lifecycles. It does not add a new transport, a new metadata channel, or a new control plane.

Before landing, the TE RDMA maintainers must agree on the hook location in submitTransferTask and WorkerPool::submitPostSend, and on the fact that the PXN packing stage runs on the submit thread while slot admission and the actual ibv_post_send remain with the existing worker pool.

Transport and PXN Responsibilities

  • The existing RDMA worker pool continues to own slice posting, CQ polling, endpoint reclaim, rail-failure accounting, and markSuccess/markFailed completion.
  • A new PxnRuntimeManager owns the following and nothing else:
    • A static topology map: for each (src_gpu, local_nic) pair, the proxy GPU that is same-rail with local_nic and reachable from src_gpu through NVLink and nvidia_peermem. The map is built once at RdmaTransport::install() and is immutable at run time.
    • A per-RDMA-context GPU slot pool allocated through ibv_reg_mr on the proxy GPU, used only as a staging buffer for in-flight cross-rail slices.
    • A per-slice packing stage (PxnPlanAndPack) that chooses between the native RDMA path and the PXN path based on peer buffer locality and source GPU affinity.
    • A memcpy dispatcher (PxnRunDispatcherOnce) that issues cudaMemcpyBatchAsync on CUDA 12.8+, or falls back to a cudaMemcpyPeerAsync loop on earlier CUDA, from the user source buffer into the proxy slot, records a CUDA event, and hands the real RDMA slice to the worker pool once the event completes.
  • A slice is admitted to the PXN path only when all of the following hold: the local source memory is GPU memory on a GPU whose PCIe root is distinct from the chosen local NIC's PCIe root, a same-rail proxy GPU exists for the chosen local NIC, that proxy GPU supports GDR and has nvidia_peermem loaded, and the peer buffer that holds the slice destination is itself registered as GPU memory on the remote side. If any condition is false the slice falls back to the native RDMA path and the controller never allocates a slot for it.
  • Fallback does not take the PXN state machine offline. The packing stage produces an explicit bypass_requests list for slices that cannot use PXN; the worker pool consumes those as ordinary RDMA requests. Topology-level unavailability (no NVLink, nvidia_peermem absent, proxy GPU without GDR, or ibv_reg_mr refusing the slot pool) disables PXN for the affected scope at startup, and the native path handles every slice from that point on.

PXN is a software staging policy. It does not change ECN or PFC on either host, and does not configure the NIC's internal congestion controls.

Data Flow and Control Hooks

Native path (unchanged)

user source buffer on NIC-local GPU
  -> WorkerPool::submitPostSend(slice_list)
  -> selectPeerDevice / enqueueSliceToOwner (per-NIC worker queue)
  -> ibv_post_send on the owner worker thread
  -> ibv_poll_cq -> markSuccess

PXN path (new, cross-rail only)

user source buffer on src_gpu
  -> PxnPlanAndPack
       - picks same-rail proxy_gpu for local_nic
       - allocates one PXN slot of size slot_size on proxy_gpu
       - creates a copy-worker (CW) shell slice that owns the slot and
         references all real RDMA slices derived from this request
  -> PxnRouteSlicesToShards: shard shells by
         (target_id, peer_device_id, src_gpu_id) % worker_count
         onto per-worker shard queues inside the WorkerPool state
  -> transferWorker: PxnRunDispatcherOnce
       - cudaMemcpyBatchAsync(src_gpu_buffer -> proxy_gpu_slot)
       - cudaEventRecord on the dispatcher stream
       - advances the CW state machine
  -> on event completion, the dispatcher pushes the real RDMA slices
     onto the WorkerPool's collective_slice_queue, rewritten to
     source = proxy_gpu_slot + offset, source_lkey = proxy_mr->lkey,
     dest_rkey = peer_mr_rkey on the same-rail peer NIC
  -> performPostSend / ibv_poll_cq -> markSuccess on the original user slice
  -> slot released once all RDMA slices derived from the CW complete

The PXN packing stage runs on the submit thread (the caller of submitTransferTask). The dispatcher and the final ibv_post_send run on the worker pool's transfer worker threads. The data path never touches new atomics or new mutexes on non-PXN slices.

Submit-side sharding of CW shells

Each CW's lifecycle, slot assignment, cascade pointer, and CUDA event are owned by exactly one worker thread, so no cross-thread coordination is needed inside a CW. The shard function is

shard_id = (target_id * 10007 + peer_device_id + src_gpu_id) % worker_count

where worker_count equals MC_WORKERS_PER_CTX (default 2). The scratch buckets used to batch per-shard pushes are thread-local on the submit thread; the shared state->pxn_slice_queue[shard_id] is protected by a per-shard TicketLock. There is no cross-submit-thread sharing of PXN state outside those shard locks.

Build Options and Environment Variables

Build

Option Default Description
USE_PXN OFF OFF strips every PXN source file, header, and pxn_info field from the build. ON compiles the PXN runtime manager, topology manager, slot pool, and WorkerPool integration into libmooncake-transfer-engine.

Compile-time OFF is the only supported way to remove PXN cost from a build. Mixing ON and OFF object files in the same translation unit set is not supported.

Runtime

Variable Default Description
MC_PXN_SLOT_SIZE_KB 1024 Size, in KB, of each GPU-memory slot on the proxy GPU. Each slot is a staging buffer for one in-flight PXN copy-worker.
MC_PXN_SLOT_COUNT_PER_WORKER 16 Number of slots assigned to each transfer worker on the proxy GPU. Total staging memory per NIC equals workers_per_ctx * slot_count_per_worker * slot_size.
MC_WITHOUT_PXN (unset) When set to 1, the PXN runtime path is bypassed even if PXN was compiled in. The topology and slot pool are still built; every slice takes the native path. Leave unset to enable PXN at run time.

Both numeric variables expect a positive integer. A non-positive or malformed value is ignored with a warning, and the built-in default is used. If a worker cannot allocate the configured staging memory on a NIC, that NIC falls back to the native path with a warning; other NICs are unaffected.

Caller Observability

PXN does not change the return layout of existing getTransferStatus or tent_task_status() queries.

  • Slices that take the PXN path report success through the same TransferStatusEnum the native path uses. The caller cannot tell from the status whether a slice was proxied; this is intentional. Transport selection is not part of the caller contract.
  • The transfer engine's existing metrics distribution (task_completion_latency_us) covers both the PXN and native paths with no new bucket layout. No PXN-specific counter is added to the public metrics surface in this change.

Configuration, Compatibility, and Rollback

  • Rollback: set MC_WITHOUT_PXN=1 at TransferEngine construction. Existing workers are not hot-switched; the current process must be drained and restarted. Compile-time USE_PXN=OFF removes every PXN field and code path, including Slice::pxn_info, which is only exposed to callers through private layout.
  • Compatibility with peers that did not compile PXN: the peer never participates in PXN packing. All PXN work happens on the sender. The peer sees an ordinary RDMA WRITE originating from the sender's same-rail NIC, which is the whole point of this proposal.
  • Compatibility with mixed-build processes: callers and the Mooncake shared library must be built with the same USE_PXN value, because Slice::pxn_info layout changes. Consumers using the public C or TENT bindings are not affected (pxn_info lives in a private slot).
  • Compatibility with other transports: PXN hooks are inside RdmaTransport and the RDMA WorkerPool only. TCP, EFA, UB, kunpeng, cxl, and TENT paths are not touched.

Performance Acceptance

Proposed gates:

Scenario Target
Compile USE_PXN=OFF Zero regression (same binary as today for non-PXN build).
Compile USE_PXN=ON, run with MC_WITHOUT_PXN=1 At most 1% regression on the native RDMA path at message sizes 1 MiB and 16 MiB.
Compile USE_PXN=ON, run with PXN enabled, source and chosen peer NIC on the same rail PXN not activated; at most 1% regression vs. native path.
Compile USE_PXN=ON, run with PXN enabled, source and chosen peer NIC on different rails Positive regression: cross-rail bandwidth should recover to within 10% of same-rail native, measured as the 1-to-1 bottleneck bandwidth.
End-to-end TTFT improvement in a representative PD-disaggregated serving benchmark (sglang bench_serving, request rate 10 req/s, max concurrency 64, 100 requests, Mean/P50/P90/P95/P99 TTFT) Measured, not asserted. In the reference deployment the P50 TTFT improved by 14.71% and the P90 by 8.61%; see the design document for the full table.

Every comparison must fix HCA, firmware, GPU, CPU/NUMA placement, GPU-to-NIC PCIe topology, message size, and workload. Each arm starts from independent process, queue, and connection state, or applies an identical documented warm-up. Soft-RoCE may be used to validate the software paths but cannot qualify physical-HCA throughput.

Before submitting a new issue...

  • Make sure you already searched for relevant issues and read the documentation
Originally created by @yunjuanyunshu0902-zyh on GitHub (Oct 23, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4897 ## Changes proposed ## Proposal Introduce a default-off cross-rail avoidance path in the classic Transfer Engine (TE) RDMA transport. When a slice originates on a GPU whose same-rail NIC is not usable for the chosen remote peer NIC, the sending host stages the payload on a locally reachable proxy GPU that shares a PCIe switch (and NVLink) with the source GPU and whose same-rail NIC can reach the remote peer NIC on the same rail. The transport issues one `cudaMemcpyAsync` from the source GPU to a pinned proxy-GPU slot, then issues the RDMA WRITE from the proxy GPU through its same-rail NIC. Receiver semantics, upper-layer transfer completion, and endpoint lifecycles remain unchanged. The goal is to remove cross-rail traffic from the spine fabric in PD-disaggregated inference. Spine switches are shared with other tenants, and cross-rail KV cache bursts observed in production trigger PFC pause-frame storms that propagate past the originating flow and degrade TTFT of co-located workloads. Keeping KV transfer traffic on the same rail eliminates the spine as a congestion domain for this workload without changing fabric configuration. This proposal is scoped to the classic Transfer Engine RDMA transport. It does not change ECN or PFC settings, does not alter the TCP, EFA, UB, or TENT transports, and does not take ownership of peer liveness or cluster membership decisions. ## Relationship to Existing Work PXN relaying was first described in the NCCL communication library (`NCCL PXN`): for a cross-rail GPU pair, egress through a locally reachable proxy GPU whose same-rail NIC can reach the remote peer on the same rail. This proposal applies the same locality principle to the Mooncake classic Transfer Engine data path. It is implemented as an opt-in build feature, keeps the existing RDMA worker pool as the submission authority, and reuses the existing endpoint, metadata, and completion lifecycles. It does not add a new transport, a new metadata channel, or a new control plane. Before landing, the TE RDMA maintainers must agree on the hook location in `submitTransferTask` and `WorkerPool::submitPostSend`, and on the fact that the PXN packing stage runs on the submit thread while slot admission and the actual `ibv_post_send` remain with the existing worker pool. ## Transport and PXN Responsibilities - The existing RDMA worker pool continues to own slice posting, CQ polling, endpoint reclaim, rail-failure accounting, and markSuccess/markFailed completion. - A new `PxnRuntimeManager` owns the following and nothing else: - A static topology map: for each `(src_gpu, local_nic)` pair, the proxy GPU that is same-rail with `local_nic` and reachable from `src_gpu` through NVLink and `nvidia_peermem`. The map is built once at `RdmaTransport::install()` and is immutable at run time. - A per-RDMA-context GPU slot pool allocated through `ibv_reg_mr` on the proxy GPU, used only as a staging buffer for in-flight cross-rail slices. - A per-slice packing stage (`PxnPlanAndPack`) that chooses between the native RDMA path and the PXN path based on peer buffer locality and source GPU affinity. - A memcpy dispatcher (`PxnRunDispatcherOnce`) that issues `cudaMemcpyBatchAsync` on CUDA 12.8+, or falls back to a `cudaMemcpyPeerAsync` loop on earlier CUDA, from the user source buffer into the proxy slot, records a CUDA event, and hands the real RDMA slice to the worker pool once the event completes. - A slice is admitted to the PXN path only when all of the following hold: the local source memory is GPU memory on a GPU whose PCIe root is distinct from the chosen local NIC's PCIe root, a same-rail proxy GPU exists for the chosen local NIC, that proxy GPU supports GDR and has `nvidia_peermem` loaded, and the peer buffer that holds the slice destination is itself registered as GPU memory on the remote side. If any condition is false the slice falls back to the native RDMA path and the controller never allocates a slot for it. - Fallback does not take the PXN state machine offline. The packing stage produces an explicit `bypass_requests` list for slices that cannot use PXN; the worker pool consumes those as ordinary RDMA requests. Topology-level unavailability (no NVLink, `nvidia_peermem` absent, proxy GPU without GDR, or `ibv_reg_mr` refusing the slot pool) disables PXN for the affected scope at startup, and the native path handles every slice from that point on. PXN is a software staging policy. It does not change ECN or PFC on either host, and does not configure the NIC's internal congestion controls. ## Data Flow and Control Hooks ### Native path (unchanged) ``` user source buffer on NIC-local GPU -> WorkerPool::submitPostSend(slice_list) -> selectPeerDevice / enqueueSliceToOwner (per-NIC worker queue) -> ibv_post_send on the owner worker thread -> ibv_poll_cq -> markSuccess ``` ### PXN path (new, cross-rail only) ``` user source buffer on src_gpu -> PxnPlanAndPack - picks same-rail proxy_gpu for local_nic - allocates one PXN slot of size slot_size on proxy_gpu - creates a copy-worker (CW) shell slice that owns the slot and references all real RDMA slices derived from this request -> PxnRouteSlicesToShards: shard shells by (target_id, peer_device_id, src_gpu_id) % worker_count onto per-worker shard queues inside the WorkerPool state -> transferWorker: PxnRunDispatcherOnce - cudaMemcpyBatchAsync(src_gpu_buffer -> proxy_gpu_slot) - cudaEventRecord on the dispatcher stream - advances the CW state machine -> on event completion, the dispatcher pushes the real RDMA slices onto the WorkerPool's collective_slice_queue, rewritten to source = proxy_gpu_slot + offset, source_lkey = proxy_mr->lkey, dest_rkey = peer_mr_rkey on the same-rail peer NIC -> performPostSend / ibv_poll_cq -> markSuccess on the original user slice -> slot released once all RDMA slices derived from the CW complete ``` The PXN packing stage runs on the submit thread (the caller of `submitTransferTask`). The dispatcher and the final `ibv_post_send` run on the worker pool's transfer worker threads. The data path never touches new atomics or new mutexes on non-PXN slices. ### Submit-side sharding of CW shells Each CW's lifecycle, slot assignment, cascade pointer, and CUDA event are owned by exactly one worker thread, so no cross-thread coordination is needed inside a CW. The shard function is ``` shard_id = (target_id * 10007 + peer_device_id + src_gpu_id) % worker_count ``` where `worker_count` equals `MC_WORKERS_PER_CTX` (default 2). The scratch buckets used to batch per-shard pushes are thread-local on the submit thread; the shared `state->pxn_slice_queue[shard_id]` is protected by a per-shard `TicketLock`. There is no cross-submit-thread sharing of PXN state outside those shard locks. ## Build Options and Environment Variables ### Build | Option | Default | Description | | --------- | -------:| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `USE_PXN` | `OFF` | `OFF` strips every PXN source file, header, and `pxn_info` field from the build. `ON` compiles the PXN runtime manager, topology manager, slot pool, and WorkerPool integration into `libmooncake-transfer-engine`. | Compile-time `OFF` is the only supported way to remove PXN cost from a build. Mixing `ON` and `OFF` object files in the same translation unit set is not supported. ### Runtime | Variable | Default | Description | | ------------------------------ | --------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `MC_PXN_SLOT_SIZE_KB` | `1024` | Size, in KB, of each GPU-memory slot on the proxy GPU. Each slot is a staging buffer for one in-flight PXN copy-worker. | | `MC_PXN_SLOT_COUNT_PER_WORKER` | `16` | Number of slots assigned to each transfer worker on the proxy GPU. Total staging memory per NIC equals `workers_per_ctx * slot_count_per_worker * slot_size`. | | `MC_WITHOUT_PXN` | *(unset)* | When set to `1`, the PXN runtime path is bypassed even if PXN was compiled in. The topology and slot pool are still built; every slice takes the native path. Leave unset to enable PXN at run time. | Both numeric variables expect a positive integer. A non-positive or malformed value is ignored with a warning, and the built-in default is used. If a worker cannot allocate the configured staging memory on a NIC, that NIC falls back to the native path with a warning; other NICs are unaffected. ## Caller Observability PXN does not change the return layout of existing `getTransferStatus` or `tent_task_status()` queries. - Slices that take the PXN path report success through the same `TransferStatusEnum` the native path uses. The caller cannot tell from the status whether a slice was proxied; this is intentional. Transport selection is not part of the caller contract. - The transfer engine's existing metrics distribution (`task_completion_latency_us`) covers both the PXN and native paths with no new bucket layout. No PXN-specific counter is added to the public metrics surface in this change. ## Configuration, Compatibility, and Rollback - Rollback: set `MC_WITHOUT_PXN=1` at TransferEngine construction. Existing workers are not hot-switched; the current process must be drained and restarted. Compile-time `USE_PXN=OFF` removes every PXN field and code path, including `Slice::pxn_info`, which is only exposed to callers through private layout. - Compatibility with peers that did not compile PXN: the peer never participates in PXN packing. All PXN work happens on the sender. The peer sees an ordinary RDMA WRITE originating from the sender's same-rail NIC, which is the whole point of this proposal. - Compatibility with mixed-build processes: callers and the Mooncake shared library must be built with the same `USE_PXN` value, because `Slice::pxn_info` layout changes. Consumers using the public C or TENT bindings are not affected (`pxn_info` lives in a private slot). - Compatibility with other transports: PXN hooks are inside `RdmaTransport` and the RDMA `WorkerPool` only. TCP, EFA, UB, kunpeng, cxl, and TENT paths are not touched. ## Performance Acceptance Proposed gates: | Scenario | Target | | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | | Compile `USE_PXN=OFF` | Zero regression (same binary as today for non-PXN build). | | Compile `USE_PXN=ON`, run with `MC_WITHOUT_PXN=1` | At most 1% regression on the native RDMA path at message sizes 1 MiB and 16 MiB. | | Compile `USE_PXN=ON`, run with PXN enabled, source and chosen peer NIC on the same rail | PXN not activated; at most 1% regression vs. native path. | | Compile `USE_PXN=ON`, run with PXN enabled, source and chosen peer NIC on different rails | Positive regression: cross-rail bandwidth should recover to within 10% of same-rail native, measured as the 1-to-1 bottleneck bandwidth. | | End-to-end TTFT improvement in a representative PD-disaggregated serving benchmark (sglang `bench_serving`, request rate 10 req/s, max concurrency 64, 100 requests, Mean/P50/P90/P95/P99 TTFT) | Measured, not asserted. In the reference deployment the P50 TTFT improved by 14.71% and the P90 by 8.61%; see the design document for the full table. | Every comparison must fix HCA, firmware, GPU, CPU/NUMA placement, GPU-to-NIC PCIe topology, message size, and workload. Each arm starts from independent process, queue, and connection state, or applies an identical documented warm-up. Soft-RoCE may be used to validate the software paths but cannot qualify physical-HCA throughput. ### Before submitting a new issue... - [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)
Author
Owner

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

Thanks for opening this issue, @yunjuanyunshu0902-zyh!

Field Value
Issue #4897
GitHub user ID 64839424
Reporter @yunjuanyunshu0902-zyh

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.

<!-- gh-comment-id:6365719372 --> @github-actions[bot] commented on GitHub (Oct 23, 2026): Thanks for opening this issue, @yunjuanyunshu0902-zyh! | Field | Value | |-------|-------| | **Issue** | #4897 | | **GitHub user ID** | `64839424` | | **Reporter** | @yunjuanyunshu0902-zyh | 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.
Sign in to join this conversation.
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
kvcache-ai/Mooncake#1244
No description provided.