═══════════════════════════════════════════════════════════════════════════════
第 0 层:Python 入口(SGLang)—— 只是普通 Mooncake 用户,对 PXN 完全无感知
═══════════════════════════════════════════════════════════════════════════════
sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
maybe_init_shared_mooncake_transfer_engine() :305
│ 条件:disaggregation_mode != null && backend == "mooncake"
▼
init_mooncake_transfer_engine() ──► MooncakeTransferEngine.__init__ :107
│
├─ TransferEngine() :122
├─ get_ib_devices_for_gpu(ib_device, gpu_id) :130 ← 本 GPU 绑哪张 HCA
└─ initialize(hostname, "P2PHANDSHAKE", protocol, device_name) :205
└─ protocol = "rdma"(默认)
── 注册 KV 显存 ── ── 发起传输 ──
register(ptr, len) :140 transfer_sync(...) :215
batch_register(ptrs, lens) :160 batch_transfer_sync(...) :236
│ │
│ engine.register_memory │ engine.batch_transfer_sync_write
▼ ▼
┌────────────────────────────────────────────────────────────────────┐
│ 注意:Python 侧没有任何 pxn 字样。PXN 是靠 C++ 侧环境变量激活的, │
│ 上层调用方式完全不变 —— 这是这个设计的关键:对上完全透明。 │
└────────────────────────────────────────────────────────────────────┘
═══════════════════════════════════════════════════════════════════════════════
第 1 层:Transport 选型(PXN 挂进来的唯一入口)
═══════════════════════════════════════════════════════════════════════════════
MC_PXN_ENABLE=1 ──► config.cpp:491 解析进 globalConfig().pxn_enable
MC_PXN_GROUP_ID :496 │ 校验:config.cpp:103-107(组名/深度/超时/rail_map)
MC_PXN_RAIL_MAP :519 │ 任一项非法 → pxn_enable 强制回落 false :525
MC_PXN_INFLIGHT_DEPTH :505 │
MC_PXN_CREDIT_TIMEOUT_MS :508 │
▼
multi_transport.cpp:361 installTransport("rdma")
if (globalConfig().pxn_enable) new PxnRdmaTransport() ← 继承 RdmaTransport
else new RdmaTransport()
═══════════════════════════════════════════════════════════════════════════════
第 2 层:install —— 建 registry / arena / 三条流水线(pxn_rdma_transport.cpp:59)
═══════════════════════════════════════════════════════════════════════════════
PxnRdmaTransport::install()
│
├─1. RdmaTransport::install(...) :62 ← 先把父类 RDMA 起好
│ 任何失败 或 !pxn_enable → 直接 return,退化成纯 RDMA
│
├─2. 收集本进程活跃 HCA getContextList() :66
├─3. cudaGetDevice(&device_id) :74 失败 → 只 WARN,退化
│
├─4. makeCudaRdmaStagingBackend(*this, dev) :83 ┐
│ │ StagingBackend 接口
│ pxn_cuda_backend.cpp:82 CudaRdmaStagingBackend│ (7 个虚函数)
│ allocateArena = cudaMalloc :89 │
│ registerArena = transport_.register :107│ ← 复用父类 MR 注册
│ exportArena = cudaIpcGetMemHandle :130│
│ importArena = cudaIpcOpenMemHandle :144│
│ ┘
├─5. LocalResources::Create(...) :84 ★ 核心初始化
│ pxn_staging.cpp:227
│ ├─ RailResolver.resolveUnique(hcas) :237 必须恰好 1 张 HCA,否则拒绝
│ ├─ Registry::Open() :243 /dev/shm/<group_id>/ + flock
│ ├─ arena_.initialize(backend) :248 ┐
│ │ pxn_staging.cpp:30 │ 1008 MB 显存
│ │ cudaMalloc → registerArena → export │ = 7 lane × 18 slot × 8 MB
│ │ ┘ (pxn_constants.h:23-31)
│ ├─ registry_->createLocal(rail, handle) :251
│ │ pxn_registry.cpp:918
│ │ /dev/shm 里建文件 → ftruncate(sizeof(ControlBlock))
│ │ → mmap → 填 header(magic/epoch/rail/arena_handle)
│ └─ registration_->publish() :256 临时名 rename 成正式名 → 别人可见
│
├─6. makeRdmaRelayBackend(*this, arena, depth) :92
│ pxn_pump.cpp:222 RdmaRelayBackend::initialize
│ selectDevice(arena) → 定位 arena 所在的 RDMA context / lkey
│ 预分配 RelaySlicePool = depth × (4096 + 8MB/slice_size) :251
│
├─7. SenderPipeline(cuda sender backend, rdma fallback, timeout) :100
├─8. RelayPipeline(control, arena, epoch, depth, relay_backend) :104
└─9. PxnPump(sender, relay) :108
→ 构造即起后台线程(pxn_pump.cpp:477)
═══════════════════════════════════════════════════════════════════════════════
第 3 层:/dev/shm 里的共享控制块 —— 同机 8 个 rank 进程互相看到的东西
═══════════════════════════════════════════════════════════════════════════════
/dev/shm/<group_id>/<uid>.<pid>.<start_ticks>.<epoch> (pxn_registry.h:104)
ControlBlock 8,306,816 B
├── RegistryHeader 448 B 本进程"我是谁、我的 arena 在哪"
│ magic / abi / epoch / heartbeat
│ rail[256] ← 我绑的是哪条 rail(如 mlx5_0)
│ arena_handle[64] ← cudaIpcMemHandle,别人靠它 map 我的显存
│
└── LaneControl lanes[7] 每条 1,186,624 B ★ 一条 lane = 一个远端 sender
├── LaneHeader 64 B
│ state kFree / kClaiming / kReady
│ sender + epoch ← 哪个进程占了这条 lane
│ arena_attached
│ head ← sender 写:已填好的最大 sequence
│ doorbell ← sender 写(GPU 或 CPU):可以搬了 ┐ 单生产
│ completed ← relay 写:已完成的最大 sequence ┘ 单消费
├── Descriptor descriptors[18] 每个 65,856 B
│ commit_sequence ← release-store 提交,兼作有效性标记
│ epoch / piece_length / plan_count / session[256]
│ plans[4096] {final_destination, length} ← 最终目的地清单
└── Completion completions[18] 每个 64 B
commit_sequence / status
┌──────────────────────────────────────────────────────────────────────┐
│ 关键点:descriptor/completion 全部 ABI 冻结(pxn_core.h:57-80 一堆 │
│ static_assert 钉死 offset 和 size),因为它们是跨进程共享内存布局, │
│ 两边编译单元必须字节一致。 │
└──────────────────────────────────────────────────────────────────────┘
显存 arena(不在 shm 里,靠 CUDA IPC 共享):
1008 MB = lane0[slot0..17] lane1[...] ... lane6[...]
slotAddress = arena + (lane_idx*18 + slot_idx) * 8MB
pxn_staging.cpp:133(本地)/ :191(导入的对端)
═══════════════════════════════════════════════════════════════════════════════
第 4 层:发送路径 —— submitTransferTask 的分流与切片
═══════════════════════════════════════════════════════════════════════════════
PxnRdmaTransport::submitTransferTask(task_list) :112
│
├─ !pxnReady() → 全部交给 RdmaTransport::submitTransferTask :114
│
├─ 逐 task 判定是否走 PXN:selectPxnLane() :209
│ ├─ opcode 必须 WRITE 且 length>0 :212
│ ├─ source 必须落在 "cuda:*" 注册段内(GPU 显存) :219-226
│ ├─ target segment 名字非空且 ≤256 :228
│ ├─ selectDevice(target) → 目标用哪个 device :236
│ ├─ canonicalize(目标 rail) :242
│ ├─ ★ 目标 rail == 本地 rail → 返回 false(不走 PXN) :244
│ │ 同 rail 直连就行,PXN 只解决"跨 rail"
│ └─ lanes_by_rail_ 缓存命中 → 直接给 lane :247
│ 否则 discover() 找同 rail 的 peer → mapPeer → addPeer :253-262
│ │
│ ├─ LocalResources::mapPeer (pxn_staging.cpp:296)
│ │ registry_->mapPeer → claimLane → attachArena → importArena
│ │ claimLane (pxn_registry.cpp:648):flock 下扫 7 条 lane
│ │ kFree → 可用;kClaiming → 清理后可用
│ │ 同 sender+epoch 已占 → 复用
│ └─ SenderPipeline::addPeer (pxn_sender.cpp:445)
│ SenderLane::Create → backend.createLane
│ pxn_cuda_backend.cpp:230
│ cudaHostRegister(control) ← shm 页锁定
│ cudaHostGetDevicePointer ← 算出 doorbell 的 GPU 地址
│ 探测 CAN_USE_STREAM_MEM_OPS:
│ 支持 → stream_mem_ops_=true(GPU 直接敲门铃)
│ 不支持 → 建 publish_event(退化成 CPU 敲门铃)
│
├─ 走 PXN 的按 (target_id, lane) 分组 :139-155
│ 每组攒 TransferSpan{source, final_destination, length, request_index}
│
├─ 逐组 buildPieces(spans, pieces) :160
│ pxn_core.cpp:90 ★ 切片规则
│ 按 8MB(kSlotSize) 和 4096(kMaxPlanCount) 双上限装箱
│ 一个 span 可跨多个 piece;一个 piece 可含多个 span
│ 失败 → 该组整体退回 direct_tasks
│
├─ 每个 piece 造一个 SenderSubmission :167-194
│ 为每个 span 分配 Transport::Slice(父类的 slice cache)
│ 挂 3 个回调:markPxnPosted / markPxnSuccess / markPxnFailed :32-42
│ → 这是与父类完成语义对接的桥:PXN 内部完成后回调它们
│ task->slice_list.push_back / total_bytes += / slice_count++
│
├─ direct_tasks 交父类 :199
├─ pending 逐个 lane->enqueue(submission) :202
└─ pump_->wake() :205
═══════════════════════════════════════════════════════════════════════════════
第 5 层:SenderLane —— 环形 ring + 信用 + fallback(pxn_sender.cpp)
═══════════════════════════════════════════════════════════════════════════════
enqueue(submission) :89
├─ validateSubmission:slices 与 spans 一一对应、长度相等 :39
├─ backend_.recordReady(fence) ← 在 cudaStreamPerThread 上打 event
│ pxn_cuda_backend.cpp:283 意义:等"计算流把 KV 算完"再搬
└─ 入 queue_(mutex 保护)
progress(made_progress) :144 ← pump 线程调
├─ pollFallbacks() :148
├─ progressFallbackQueue() :153 fallback 队列非空则不发新的
├─ 若有 pending_publication_(CPU 门铃路径): :158
│ pollCpuDoorbell → ready 则 finishPublication(cpu_doorbell=true)
├─ 取队首,判定是否 fallback: :186-195
│ quarantined / next_sequence_==0 → fallback
│ !hasRingCredit(next, reaped) 且等待超过 credit_timeout → fallback
│ hasRingCredit: next-reaped <= 18 (pxn_core.cpp:76)
└─ publish(submission, next_sequence_, ...) :210
│
├─ slotIndex(seq) = (seq-1) % 18 pxn_core.cpp:71
├─ prepareDescriptor(piece, session, epoch, desc) :360
│ pxn_core.cpp:122 从 epoch 字段起清零,填 plans[]
│ ★ 注意:不含 commit_sequence —— 它最后单独提交
├─ 算 copies:piece 内各 span 在 slot 里连续排布 :378-392
│ destination = arena + slot_offset,逐 span 累加
├─ commitDescriptor(desc, seq) :393
│ release-store commit_sequence ← 发布屏障
└─ backend_.publish(handle, fence, copies, seq, state)
pxn_cuda_backend.cpp:326
├─ cudaStreamWaitEvent(lane->stream_, ready_event) 等计算完
├─ cudaMemcpyBatchAsync(dst[], src[], sizes[]) :363/369
│ CUDART>=12080 用批量 API;否则退化成逐个 cudaMemcpyAsync :378
│ ★ 这是 commit5 优化的 piece→slice 路径
└─ 敲门铃:
stream_mem_ops_ → cuStreamWriteValue64(doorbell, seq) :391
state = kPublished(GPU 自己敲,零 CPU 介入)
否则 → cudaEventRecord(publish_event) :400
state = kPendingCpuDoorbell(等 CPU 补敲)
finishPublication() :418
├─ head = sequence (release)
├─ cpu_doorbell ? doorbell = sequence (release)
├─ markPosted(slices) → 父类 Slice::POSTED
└─ inflight_.push_back({seq, slices}); next_sequence_ = seq+1
reap(made_progress) :110 ← pump 线程调
└─ 逐 inflight_ 队首:读 completions[slot]
completion_sequence != published.sequence → break(还没完成)
tryLoadCompletion → status==0 ? markSuccess : markFailed
reaped_sequence_ = seq(释放 ring 信用)
── Fallback 路径(PXN 走不通时保底)── :582
RdmaSenderFallback::submit
把 piece.spans 直接翻成父类 TransferRequest{source → final_destination}
allocateBatchID → RdmaTransport::submitTransferTask :622
RdmaFallbackTransfer::poll 逐 index 查状态 → mark_success/failed :535
★ 语义:绕过 relay,本 rank 自己的 HCA 直发,慢但一定能到
═══════════════════════════════════════════════════════════════════════════════
第 6 层:RelayPipeline —— arena 持有者(被代理方)把数据转发出去
═══════════════════════════════════════════════════════════════════════════════
progressInbound() pxn_pump.cpp:402
├─ inflight_count_ >= max_inflight_ → 退出(背压)
├─ 轮询 7 条 lane(next_lane_ 轮转,保证公平) :406
│ lane state 必须 kReady
│ doorbell < next_sequence_ → 跳过(sender 还没敲) :417
├─ validateDescriptor(desc, seq, epoch) :422
│ pxn_core.cpp:156 校 sequence/epoch/长度/plan 累加和
│ 不合法 → 塞 CompletedRelayTransfer(ERR_INVALID_ARGUMENT),不中断整条 lane
├─ 组 RelaySubmission{source=slot 地址, session, plans} :434
└─ backend_->submit(submission, transfer)
RdmaRelayBackend::submit (pxn_pump.cpp:147)
├─ forEachSlice 两遍:先数 slice 数,再填 :150/162
│ forEachSlice (:323) 逐 plan 按 3 者取小切分:
│ min(remaining, slice_size, route->end - destination)
├─ resolveRoute(session, destination) :258
│ session → target_id → SegmentDesc(缓存在 sessions_)
│ selectDeviceByLocalHca / selectDevice(看 affinity 配置)
│ 缓存 TargetRoute{begin, end, rkey, device_id, peer_nic_path}
├─ pool_->acquire(count) 不够 → BatchBusy(回落 CompletedRelayTransfer)
├─ 填 Slice:source=arena slot, dest_addr, lkey=arena, rkey=target
│ completion_callback = RdmaRelayTransfer::completeSlice :177
└─ source_context_->submitPreparedPostSend(slices) :193
★ 直接 post 到 arena 所在 HCA 的 QP,绕过父类 slice 调度
reapInbound() :374
└─ 逐 lane 队首 poll:remaining_==0 才算完成 :112
complete(lane, seq, status) :468
commitCompletion(completions[slot], seq, status) release-store
header.completed = seq release-store
↓
sender 侧 reap() 就能看到了
═══════════════════════════════════════════════════════════════════════════════
第 7 层:PxnPump —— 单线程驱动全部四个阶段(pxn_pump.cpp:489)
═══════════════════════════════════════════════════════════════════════════════
PxnPump::run() while(running_)
┌──────────────────────────────────────────────────┐
│ 1. sender_.reapOutbound() 收自己发出去的完成 │
│ 2. relay_.reapInbound() 收自己代转的完成 │
│ 3. relay_.progressInbound() 代转新的 │
│ 4. sender_.progressOutbound() 发新的(lane 轮转) │
└──────────────────────────────────────────────────┘
四步都没进展 → condition_.wait_for(50us) ← 忙等但有退让
wake() 由 submitTransferTask 末尾触发 :205
★ 一个进程同时是 sender(用别人的 arena)和 relay(借出自己的 arena),
所以同一个 pump 线程要同时驱动两个方向。
═══════════════════════════════════════════════════════════════════════════════
端到端:一次跨 rail 的 KV 写,完整时序
═══════════════════════════════════════════════════════════════════════════════
进程 A (GPU0, rail=mlx5_0) 进程 B (GPU1, rail=mlx5_1)
想写到远端节点,但目标在 mlx5_1 那条 rail 持有 mlx5_1,出借 arena+lane
───────────────────────────────── ──────────────────────────────
Python batch_transfer_sync_write
▼
submitTransferTask :112
▼
selectPxnLane → 目标 rail=mlx5_1
≠ 本地 mlx5_0 → 走 PXN :244
▼
discover() 找到 B ────────────► /dev/shm/<group>/B 的 ControlBlock
claimLane → lane_k lanes[k].state: kFree→kReady
attachArena lanes[k].arena_attached = 1
importArena(B.arena_handle) ← CUDA IPC,A 拿到 B 显存的地址
▼
buildPieces:按 8MB / 4096 装箱 :160
▼
为每 span 分配父类 Slice + 3 回调 :174
▼
lane->enqueue → recordReady(event) 在计算流打 event
▼
═══ pump 线程 ═══
progress() → hasRingCredit?
无信用且超时 → ┄┄► RDMA fallback 直发(绕过 B)
有信用 ↓
prepareDescriptor → 填 plans[] ──► lanes[k].descriptors[slot]
commitDescriptor(seq) release ──► commit_sequence = seq
▼
cudaStreamWaitEvent(计算完)
cudaMemcpyBatchAsync ──► ★ NVLink/PCIe: A显存 → B的arena slot
▼ (8MB slot,数据现在在 B 手里)
cuStreamWriteValue64(doorbell) ──► lanes[k].header.doorbell = seq
(或 CPU 补敲 via publish_event) ▲
markPosted → Slice::POSTED │
inflight_.push({seq, slices}) │
│
═══ B 的 pump 线程 ═══
progressInbound() :402
doorbell >= next_seq ✓
validateDescriptor :422
▼
RelaySubmission
source = arena+slot
plans = 最终目的地清单
▼
forEachSlice:按 slice_size 和
route 边界切分 :323
resolveRoute → rkey/nic_path :258
▼
submitPreparedPostSend
★ 从 B 的 mlx5_1 出网 ═══════► 远端节点
▼
reapInbound():remaining_==0
commitCompletion(seq, st) :472
header.completed = seq
│
┌─────────────────────────────────┘
▼
═══ A 的 pump 线程 ═══
reap() 读 completions[slot] :110
seq 匹配 → status==0 ?
markSuccess → Slice::markSuccess ──► 父类 task 完成计数
markFailed → Slice::markFailed
reaped_sequence_ = seq ← ring 信用归还
▼
Python 侧 batch_transfer_sync_write 返回
┌────────────────────────────────────────────────────────────────────────┐
│ 一句话:A 没有 mlx5_1 的直连能力,于是把 KV 通过 NVLink 搬进 B 的显存 │
│ 中转区,敲个门铃;B 的 pump 看见门铃,用自己的 HCA 把数据 RDMA 发出去,│
│ 完成后回写 completion;A 轮询到 completion 才认为这次写成功。 │
│ 全程 A 的 Slice 完成语义不变,所以对父类 RdmaTransport 和 Python 透明。│
└────────────────────────────────────────────────────────────────────────┘
═══════════════════════════════════════════════════════════════════════════════
退化路径汇总(每一层都能独立回落,任一环失败不影响正确性)
═══════════════════════════════════════════════════════════════════════════════
pxn_enable=0 / 配置非法 → 装 RdmaTransport multi_transport:361
cudaGetDevice 失败 → install 提前返回,纯 RDMA pxn_rdma:74
本地 HCA 不唯一 → LocalResources::Create 拒绝 pxn_staging:238
registry/arena 建不起来 → WARN + 纯 RDMA pxn_rdma:87/95
非 WRITE / 非 GPU 源 / 同rail→ direct_tasks 交父类 pxn_rdma:212-244
找不到同 rail 的 peer → direct_tasks pxn_rdma:264
buildPieces 失败 → 整组退回 direct_tasks pxn_rdma:161
ring 无信用且超 credit 超时 → RDMA fallback 直发 pxn_sender:190
recordReady 失败 / fence 空 → RDMA fallback pxn_sender:202
descriptor 准备/提交失败 → RDMA fallback + quarantine pxn_sender:363/393
GPU 不支持 stream mem ops → CPU 门铃(publish_event) pxn_cuda:270
CUDART < 12080 → 逐个 cudaMemcpyAsync pxn_cuda:378
relay slice pool 耗尽 → BatchBusy → CompletedRelayTransfer(err)
descriptor 校验失败 → 该 seq 记错误,lane 继续 pxn_pump:424