xserv/docs/18-pipeline-parallelism.md

Phase 18: Pipeline Parallelism (PP)

目标：在单机多卡上做 流水线并行，把 Qwen3-8B 的 层 切成 P 段（stage）， 每张卡只持有连续的一段层（+ stage0 的 embed_tokens、最后一段的 norm/lm_head）， 激活（hidden state）在相邻 stage 之间用 NCCL P2P send/recv 传递。 与 TP（按 head / 中间维切，每层 2 次 AllReduce）互补：PP 通信量小（每 token 仅 P-1 次点对点传 [tokens, hidden]），KV 与权重按 层 降到约 1/P。 先做 PP=2 / 4（组内），正确性优先。

1. 硬件约束（dash5）

• 8× RTX 5090（32GB，SM120），无 NVLink，纯 PCIe Gen5。
• 拓扑：GPU 0–3 一组、4–7 一组，组内 PHB（同 host bridge，可 P2P），跨组 NODE。
• PP 同样建议在组内（0–3 或 4–7）：虽然 PP 的通信量远小于 TP，但 P2P 仍走 PCIe， 跨组延迟更高。PP=2/4 用 0–1 / 0–3。
• 相比 TP：TP 每 token 2·layers = 72 次 AllReduce（延迟主导）；PP 每 token 仅 P-1 次 send/recv，每次 [tokens, hidden] BF16（decode batch=1 时 8KB）。 PP 对慢互联（PCIe / 无 NVLink）更友好，这是在 dash5 上做 PP 的主要动机之一。

2. 切分方案（layer-wise）

Qwen3-8B：hidden=4096、num_heads=32、num_kv_heads=8、head_dim=128、 intermediate=12288、layers=36、vocab=151936。36 能被 2/4 整除（PP=3/6 需处理余数， 本阶段先要求 layers % P == 0）。

设 stage 数 P，本 stage = s，每段 L = layers / P 层，本段持有全局层 [s·L, (s+1)·L)：

组件	持有者	说明
embed_tokens [vocab, hidden]	仅 stage 0	token → hidden
transformer block i 的全部权重	持有 i 的那个 stage	不切 head / 中间维（与 TP 正交）
该层 KV cache	持有 i 的那个 stage	每卡 KV 降到约 1/P
最终 norm [hidden]	仅最后一段
lm_head [vocab, hidden]	仅最后一段	hidden → logits

• 注意力 / MLP 的层内计算 完全不变（不需要 AllReduce）：每个 stage 用它自己那几层 的完整权重、完整 head 做 forward。PP 与 TP 正交，可叠加（本阶段不实现 TP×PP）。
• RoPE 用全局绝对 position，每个 stage 的 RopeCache 完全相同（按 position 索引）， 各 stage 独立做，无需通信。
• 每个 stage 一个独立的 PagedKVCache，层数 = 本段层数 L（不是 36）。forward 时 按「本段内的局部层号 0..L」索引 cache —— 与单卡代码完全一致，只是 self.layers 只装了本段的层。实现技巧：给 cache 传一个 num_hidden_layers 改写成 L 的 config 克隆， 无需改 PagedKVCache。

通信点

• prefill：stage s 算完本段层，得到 [S, hidden] → send 给 s+1；s+1 recv 后接着算。
• decode：同理传 [B, hidden]（batch=1 时 [1, hidden]）。
• 每 token 共 P-1 次 send/recv；最后一段算出 logits 并采样。
• 采样得到的 token id（一个 u32）由 最后一段经线程内 channel 回传给 stage0 （同进程多线程，无需走 NCCL）。

3. 进程 / 线程模型

沿用 TP 的 单进程、多线程：每个 stage 一个 OS 线程，线程启动时 cudaSetDevice(stage)。

• stage 0 = 协调者（coordinator），跑在调用线程上：持有 scheduler、tokenizer、HTTP response sender、停止判定（eos / max_tokens）与「下一步输入 token」。
• stage 1..P-1 = worker 线程：从控制 channel 收命令（Register/Prefill/Decode/Free/Shutdown）， 每步 recv 上游 hidden → 跑本段层 → send 给下游；最后一段 head+采样 → 把 token 回传 stage0。
• 控制信息（命令、采样参数、token id）走 mpsc（极小）；重活（hidden 张量）走 NCCL P2P（GPU↔GPU）。

v1 串行语义：一次处理一个请求、一次一个 token，流水线每步「灌满又排空」 （stage0 decode 第 t+1 步依赖最后一段第 t 步采出的 token）。这保证 正确性， 并拿到 TTFT/TPOT 与每卡显存；throughput 的真正收益来自 microbatch/请求级流水线 重叠（1F1B），列为后续工作（见 §7）。

执行流（每请求）：

coordinator                     worker s (1..P-1)            last stage (P-1)
─────────────                   ─────────────────            ────────────────
broadcast Register(slot)        cache.register(slot)         cache.register(slot)
broadcast Prefill{n,slot,samp}
 x=embed(prompt)
 x=layers_prefill(x,slot)
 send x → stage1                recv x ← s-1
                                x=layers_prefill(x,slot)
                                send x → s+1 ───────────────► recv x ← P-2
                                                             x=layers_prefill(x,slot)
                                                             logits=head(x); next=sample
 next ◄────────────── token channel ◄──────────────────────  token_tx.send(next)
 stream(next); loop Decode{slot} 直到 eos/length
broadcast Free(slot)            cache.free(slot)             cache.free(slot)

4. 通信库：NCCL P2P

复用 xserv-distributed（已有 NCCL FFI + TpContext/AllReduce），新增：

• FFI：ncclSend(sendbuff, count, dtype, peer, comm, stream)、 ncclRecv(recvbuff, count, dtype, peer, comm, stream)。
• PpContext：与 TpContext 同样的 ncclCommInitRank（一个 comm 跨 P 个 stage）， 外加 send_bf16_ptr(ptr, count, peer) / recv_bf16_ptr(ptr, count, peer)，在 null stream 上发起（与模型 kernel 同流，天然有序）。
• 线性流水线无死锁：stage0 只 send、最后一段只 recv、中间段「先 recv 上游、再 send 下游」， 依赖链无环，从头解锁。每个 stage 在 send/recv + 本段计算后 synchronize()， 确保 NCCL 读完发送缓冲再复用/释放（v1 串行下成本可接受）。

决策点：和 TP 一样，collective/P2P 先用 NCCL 把 PP 跑通拿正确性与基线； 手写 P2P（PCIe 上的 cudaMemcpyPeer）作为后续学习项。

5. 权重分片加载

Qwen3::from_weights_pp(config, weights, stage, num_stages, device)：

• 只把全局层 [s·L, (s+1)·L) 搬到本 stage 的 GPU（其余层的权重直接 drop，不占显存）。
• embed_tokens：仅 stage 0 加载；其余 stage 放一个 1×1 占位张量（forward 用 is_first_stage 守卫，永不触碰）。
• norm/lm_head：仅最后一段加载；其余放占位。
• head 不切（不做 TP），所以 local_num_heads = num_heads、local_num_kv_heads = num_kv_heads。

每卡显存 ≈ 权重(transformer 1/P) + KV(1/P) + (stage0: embed) + (last: norm+lm_head)。 对 Qwen3-8B：transformer 层约 14GB，PP=2 每卡约 7GB 层权重 + embed 或 lm_head（各 ~1.2GB）。

6. 实现步骤（逐步可验证）

1. P18.1 — xserv-distributed P2P：ncclSend/Recv FFI + PpContext。 验收：2 卡，rank0 send 已知向量、rank1 recv，校验一致（tests/sendrecv.rs）。
2. P18.2 — 分段权重加载：from_weights_pp，每 stage 只持有本段层 + 该有的 embed/head。 验收：各 stage 层数 = L、显存约 1/P（+ embed/head）。
3. P18.3 — stage forward：embed / forward_layers_prefill / forward_layers_decode / head，每段独立 KV cache。 验收：PP=1 与单卡 forward_*_paged 逐 token 一致（同一条代码路径退化）。
4. P18.4 — PP engine + --pp N：多线程 stage workers + NCCL 传递 + stage0 协调。 验收：--pp 2/4 端到端可服务；greedy 输出与单卡（PP=1）逐 token 一致； 用现有 llama.cpp bench 跑正确性（GSM8K/AIME）；测 PP=1/2/4 的 TTFT/TPOT/每卡显存。

7. 预期与风险

• 显存：每卡 transformer 权重 + KV ≈ 1/P，这是 PP 的主要收益（可上更大模型 / 更长 context）。
• 单流吞吐：v1 串行无 stage 重叠 → 单流 tok/s 不会超过单卡（多一份 P2P + sync 开销， 可能略低）。这是 PP 的本质：没有 microbatch 重叠就没有加速。诚实记录实测，并与 llama.cpp 的 --split-mode layer（同样是层切流水线、单序列也串行跨卡）对比 —— 两者单流 都应≈单卡。
• 真正的 throughput 收益（后续）：请求级 / microbatch 流水线（1F1B），让 stage 间重叠： stage1 算 microbatch A 时 stage0 算 B。需要把 scheduler 改成跨 stage 连续批处理。
• 风险：NCCL 多线程 init 同步；send 缓冲生命周期（必须 sync 后再复用）； layers % P != 0 的余数分配（本阶段先约束整除）；与 CUDA Graph decode 的结合（先走非 graph 路径）。
• 正确性优先：先 PP=1 等价（逐 token 对齐），PP=2/4 与单卡对齐，再谈性能。

8. 与 llama.cpp 的对比口径

• xserv：--pp N，CUDA_VISIBLE_DEVICES=0..N-1。
• llama.cpp：-sm layer（默认即层切流水线）+ --tensor-split 均分层，CUDA_VISIBLE_DEVICES=0..N-1。 （对照 TP 用的是 -sm row。）
• 指标：正确性（GSM8K / AIME exact-match）、单流 TTFT/TPOT、并发吞吐、每卡 VRAM。
• 复用 tools/bench/runner.py 与 run_pp_parallel.sh（仿 run_tp_parallel.sh）。

9. 不在本阶段范围

• TP×PP 混合（2D 并行）、跨组 / 多节点。
• microbatch / 1F1B 流水线重叠（throughput 收益，后续）。
• vocab-parallel embedding / lm_head。
• layers % P != 0 的非均匀切分；与 CUDA Graph decode 结合。