xserv/docs/17-tensor-parallelism.md

Phase 17: Tensor Parallelism (TP)

目标：在单机多卡上做 张量并行，把 Qwen3-8B 的权重、计算和 KV cache 按 head / 中间维切分到 TP 个 GPU 上，用 AllReduce 聚合，降低单卡显存压力并提升吞吐。 先做 TP=2 / 4（组内），跳过投机解码（原 Phase 16）。

1. 硬件约束（dash5）

• 8× RTX 5090（32GB，SM120），无 NVLink，纯 PCIe Gen5。
• 拓扑：GPU 0–3 一组、4–7 一组，组内 PHB（同 host bridge，可 P2P），跨组 NODE。
• TP 必须在组内（0–3 或 4–7），否则 AllReduce 走跨组 PCIe，延迟更高。
• AllReduce 带宽受限于 PCIe（~单向 64GB/s），远低于 NVLink；通信会是 decode 的主要开销。

2. 切分方案（Megatron-style）

Qwen3-8B：hidden=4096、num_heads=32、num_kv_heads=8、head_dim=128、 intermediate=12288、layers=36、vocab=151936。TP=2/4/8 都能整除 32/8/12288。

每个 transformer block 的切分（设 world size = T，本 rank = r）：

Attention（column → row）

权重	原 shape (已转置)	切分	每 rank shape
q_proj_wt	[hidden, num_heads·head_dim]	column（按 Q head）	[hidden, (num_heads/T)·head_dim]
k_proj_wt	[hidden, num_kv_heads·head_dim]	column（按 KV head）	[hidden, (num_kv_heads/T)·head_dim]
v_proj_wt	同上	column	同上
o_proj_wt	[num_heads·head_dim, hidden]	row	[(num_heads/T)·head_dim, hidden]

• 每个 rank 只算自己的 num_heads/T 个 Q head 和对应的 num_kv_heads/T 个 KV head； GQA 的 n_rep = num_heads/num_kv_heads = 4 在每个 rank 内保持不变。
• q_norm/k_norm（[head_dim]）逐 head 应用，复制到每个 rank。
• RoPE 逐 head、按 position 应用，每个 rank 独立做。
• KV cache 也切分：每个 rank 的 paged KV 只存自己的 num_kv_heads/T 个 head → 每卡 KV 显存降为 1/T（TP 的一大收益）。
• attn = merge_heads(...) @ o_proj_wt 得到部分 [T_tok, hidden] → AllReduce(sum) → 完整。

MLP / SwiGLU（column → row）

权重	原 shape	切分	每 rank shape
gate_proj_wt	[hidden, intermediate]	column	[hidden, intermediate/T]
up_proj_wt	同上	column	同上
down_proj_wt	[intermediate, hidden]	row	[intermediate/T, hidden]

• silu(gate)*up 在切分后的 [T_tok, intermediate/T] 上逐元素做，无需通信。
• down = (...) @ down_proj_wt 得到部分 [T_tok, hidden] → AllReduce(sum) → 完整。

复制（不切分）

• 所有 RMSNorm 权重（input_norm/post_norm/最终 norm）：每个 rank 在 AllReduce 后 拿到完整 hidden，本地用复制的权重归一化。
• 第一版：embed_tokens 和 lm_head 复制（各 ~1.2GB）。 后续优化：vocab-parallel embedding（local lookup + AllReduce）、column-parallel lm_head + AllGather。

通信点

每层 2 次 AllReduce（o_proj 后、down_proj 后）→ 每生成 1 token 共 2·36 = 72 次。 decode 时每次 AllReduce 张量是 [batch, 4096] BF16（单 token batch=1 时 8KB），延迟主导。

3. 进程 / 线程模型

单进程、多线程：每个 TP rank 一个 OS 线程，线程启动时 cudaSetDevice(rank_device) 并绑定。

选择理由：

• xserv 的 caching allocator 是 thread_local，每线程独立池 → 天然契合「一线程一卡一池」。
• CUDA context 隐式按 device/thread 管理；线程内只 set 一次 device、不再切换即可。
• HTTP server / 调度器仍在主线程（rank 0 协调），无需多进程 IPC，改动最小。
• 单机 8 卡足够；多进程（torchrun 式）留待真正跨节点时再说。

执行流：主线程调度器准备一个 step 的输入（tokens/positions/slots），广播给 T 个 rank 线程； 每个 rank 线程跑自己的分片 forward（含层内 AllReduce），rank 0 拿到完整 logits 后采样。 用 barrier / channel 同步每个 step。

4. 通信库：NCCL

用 NCCL（dash5 已装：/usr/lib/x86_64-linux-gnu/libnccl.so.2，/usr/include/nccl.h）。

• 新建 crate xserv-distributed：NCCL FFI（ncclGetUniqueId、ncclCommInitRank、 ncclAllReduce、ncclGroupStart/End）+ TpContext（rank/world/comm/stream）+ all_reduce_sum(&mut GpuBuffer) 原语。
• NCCL 多线程模式：主线程生成 ncclUniqueId，各 rank 线程用 ncclCommInitRank(comm, world, id, rank) 初始化（需 ncclGroupStart/End 包裹并发 init）。
• AllReduce 用 BF16（ncclBfloat16）+ ncclSum，在每个 rank 自己的 stream 上。

决策点：collective 用 NCCL 还是自己手写 P2P ring/tree AllReduce？ 本项目是「从零构建」，但 collective 属于基础设施（类比我们也用 cuBLAS 作为可切换后端）。 推荐先用 NCCL 把 TP 跑通、拿到正确性与加速比，后续可选做手写 AllReduce 作为学习项。

5. 权重分片加载

每个 rank 只加载/保留自己的分片，省显存：

• column-parallel 权重：按输出维切片取本 rank 的 [*, dim/T] 段。
• row-parallel 权重：按输入维切片取本 rank 的 [dim/T, *] 段。
• 复制权重（norm/embed/lm_head）：每个 rank 各留一份。

实现：loader 读 safetensors（mmap）时按 rank 只搬运需要的切片到该 rank 的 GPU； Qwen3::from_weights_tp(config, weights, rank, world) 在转置/切分时按 rank 取段。

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

1. P17.1 — xserv-distributed 基础：NCCL FFI + TpContext + all_reduce_sum。 验收：2 卡各放一个已知向量，AllReduce 后两卡结果都等于和。
2. P17.2 — 分片权重加载：from_weights_tp，每 rank 只持有自己的分片。 验收：各 rank 权重 shape 正确、显存占用约为 1/T（+ 复制项）。
3. P17.3 — TP forward：rank 内 attention/MLP + 层内 AllReduce。 验收：TP=2 的 logits 与 TP=1 在 BF16 容差内一致（top-1 一致，top-5 重合）。
4. P17.4 — 接入 engine/server：--tp N，多线程 rank workers + rank0 调度。 验收：--tp 2 端到端可服务；用现有 llama.cpp bench 跑正确性； 测 TP=2 vs TP=1 的吞吐 / TTFT / 每卡显存。

7. 预期与风险

• 显存：每卡权重 + KV 降到约 1/T（embed/lm_head 暂复制）。TP=2 时单卡 ~8GB 权重 + 更大 KV/并发空间。
• 吞吐：PCIe AllReduce 延迟会吃掉部分收益；decode 是延迟敏感的，72 次小 AllReduce/token 可能让 TP=2 的单流 tok/s 不一定线性提升，但能跑更大 batch / 更长 context。先测实测数。
• 风险：NCCL 多线程初始化的同步、每 rank stream 与现有 kernel stream 的协调、 KV cache 按 rank 切 head 后 paged kernel 的 head 维参数要用 per-rank 值。
• 正确性优先：先 TP=1 等价（logits 对齐），再谈性能。

8. 不在本阶段范围

• 跨组 TP=8、Pipeline Parallelism、多节点。
• vocab-parallel embedding / lm_head（先复制）。
• 手写 AllReduce（NCCL 跑通后可选）。
• 与 CUDA Graph decode 的结合（先走非 graph 路径）。