第 9 章 多进程与张量并行

在前面的概念章节中，我们已经明确区分了 Multi-process 与 Tensor Parallelism 的职责边界：Multi-process 负责组织多个进程协同工作，而 Tensor Parallelism 则将同一层的模型权重与计算切分到多个 rank 上执行。本章将深入 mini-SGLang 的源码实现，通过完整的代码路径将这两条技术线串联起来。

建议先阅读概念文档理解原理，再对照本章代码的实现。

1 本章学习目标

完成本章学习后，应该能够回答以下问题：

• 区分 ZMQ 请求广播、Gloo 控制通信与 NCCL 张量通信的区别与联系。
• 解释 Column/Row Parallel 的权重维度、输入输出形状及通信位置。
• 说明 QKV、词表和 KV Cache 的分片规则，以及 KV head 复制与词表整除的实现边界。
• 从接收请求追踪到生成下一 token，说明各 rank 如何维持相同的请求顺序与计算状态。

2 定位源码

本章的主要代码文件位于 python/minisgl/：

文件	作用
server/launch.py	解析 TP 等参数，启动 API Server，并创建 Scheduler、Tokenizer、Detokenizer 子进程，等待所有进程就绪
distributed/info.py	保存当前进程的rank、size，供各层通过 get_tp_info() 查询 TP 信息
distributed/impl.py	用统一接口封装 Torch Distributed 与 PyNCCL 的 All-Reduce、All-Gather
engine/engine.py	绑定 GPU（cuda:{rank}）、初始化通信、加载并切分模型权重、分配 KV Cache、执行前向与采样
scheduler/io.py	只有 rank 0 负责与 Tokenizer/Detokenizer 交互，向其他 rank 广播请求并同步本轮消息数量
scheduler/scheduler.py	主调度循环：更新请求状态、选择 Prefill/Decode batch、准备输入、调用 Engine 前向
scheduler/decode.py	管理 running decode 请求；schedule_next_batch 时按 uid 排序，保证多 rank 间 batch 顺序一致
kernel/pynccl.py	rank 0 创建并广播 NCCL Unique ID，初始化 PyNCCL Communicator
kernel/csrc/src/pynccl.cu	通过 TVM FFI 调用 NCCL，在 CUDA stream 上提交 all-reduce、all-gather 通信
models/weight.py	按参数名对 checkpoint 进行切分Column/Row/Vocab，并合并 QKV、gate/up 等局部投影分片

3 从启动到输出结果

接下来，我们会跟着一个请求，完整走一遍从系统启动到最终输出结果的全流程，厘清各个组件之间是如何协作的。

---

3.1 创建进程

命令行启动推理模型时，需要根据模型总参数量与显卡容量确定 tensor parallel 的规模（tp_size）。例如使用 3 张 RTX 5090 运行 Qwen3-32B：

python -m minisgl --model "Qwen/Qwen3-32B" --tp-size 3

注： 如果不指定 --tensor-parallel-size/--tp-size，默认为单卡运行。

在 server/args.py 的 parse_args 中，会先构造：

kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"])

即父进程解析参数后，server_args.tp_info 初始为 DistributedInfo(rank=0, size=tp_size)。

接下来，launch_server 解析完参数后，会把 start_subprocess 作为回调传给 run_api_server：

# server/launch.py（节选）
def start_subprocess() -> None:
    mp.set_start_method("spawn", force=True)
    world_size = server_args.tp_info.size
    ack_queue: mp.Queue[str] = mp.Queue()

    # 启动各 rank 的 scheduler
    for i in range(world_size):
        new_args = replace(
            server_args,
            tp_info=DistributedInfo(i, world_size),
        )
        mp.Process(
            target=_run_scheduler,
            args=(new_args, ack_queue),
            daemon=False,
            name=f"minisgl-TP{i}-scheduler",
        ).start()

    # 再启动 1 个 detokenizer + num_tokenizers 个 tokenizer
    ...
    # 等待确认（只有 rank0 的 scheduler 会发 ack）
    # 总共期望：1 (primary scheduler) + 1 (detokenizer) + num_tokenizers
    for _ in range(num_tokenizers + 2):
        logger.info(ack_queue.get())

run_api_server(server_args, start_subprocess, run_shell=run_shell)

父进程为每个 rank i 使用 dataclasses.replace 将 tp_info 替换为 DistributedInfo(i, world_size)，以 spawn 方式启动对应的 scheduler 子进程。

通过 ack_queue 等待确认 num_tokenizers + 2 条信号后，start_subprocess 才返回调用。

# server/api_server.py（节选）
def run_api_server(config, start_backend, run_shell) -> None:
    global _GLOBAL_STATE
    ...
    host = config.server_host
    port = config.server_port

    _GLOBAL_STATE = FrontendManager(
        config=config,
        recv_tokenizer=ZmqAsyncPullQueue(...),
        send_tokenizer=ZmqAsyncPushQueue(...),
    )

    # 启动 backend，同时启动各个 rank 的子进程
    # 这里调用的就是上面的 start_subprocess
    start_backend()   
    logger.info(f"API server is ready to serve on {host}:{port}")
    if not run_shell:
        uvicorn.run(app, host=host, port=port)
    else:
        asyncio.run(shell())

run_api_server 先建立前端 ZMQ 通道，再调用 start_backend（即 start_subprocess）创建各个子进程；等该回调返回后才进入 uvicorn.run ，而“进程已创建”与“HTTP 可接收请求”之间还有模型与通信初始化。

---

3.2 初始化 Engine 与通信

Scheduler.init 首先创建 Engine(config)，再依次初始化 Table、Cache、Decode、Prefill 管理器，通过 SchedulerIOMixin 初始化 IO：

# scheduler/scheduler.py（节选）
self.engine = Engine(config)
...
self.table_manager = TableManager(...)
self.cache_manager = CacheManager(...)
self.decode_manager = DecodeManager(...)
self.prefill_manager = PrefillManager(...)
...
super().__init__(config, self.engine.tp_cpu_group)  # IO mixin

Engine 的关键就是初始化各个组件：

# engine/engine.py（节选）
set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)  # 必须最先设置
self.device = torch.device(f"cuda:{config.tp_info.rank}")
torch.cuda.set_device(self.device)
self.stream = torch.cuda.Stream()
torch.cuda.set_stream(self.stream)

self.tp_cpu_group = self._init_communication(config)

# 初始化模型，先在 meta device 上建结构，再加载权重
with torch.device("meta"), torch_dtype(config.dtype):
    self.model = create_model(config.model_config)
self.model.load_state_dict(self._load_weight_state_dict(config))

顺序为：

• 先设置 TP 信息，再创建模型层，各层通过 get_tp_info() 计算本地权重形状。
• 设备编号是可见设备中的编号，直接用 cuda:{config.tp_info.rank}。
• 先在 meta device 上建结构，再接入权重。

其中通信初始化有两条路径，默认走 PyNCCL ：

# engine/engine.py（节选）
def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
    # size==1 或默认 use_pynccl 都走 Gloo 路径；
    # 不创建 PyNCCL 的 return 在 enable_pynccl_distributed 内部
    if config.tp_info.size == 1 or config.use_pynccl:
        torch.distributed.init_process_group(
            backend="gloo",
            rank=config.tp_info.rank,
            world_size=config.tp_info.size,
            ...
        )
        tp_cpu_group = torch.distributed.group.WORLD
        max_bytes = (...)
        enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes)
    else:
        torch.distributed.init_process_group(
            backend="nccl",
            rank=config.tp_info.rank,
            world_size=config.tp_info.size,
            ...
        )
        tp_cpu_group = torch.distributed.new_group(backend="gloo")
    return tp_cpu_group

配置	PyTorch 默认进程组	GPU 集合通信	CPU 控制通信
tp_size=1	Gloo	无	Gloo WORLD
tp_size>1（默认）	Gloo	独立 PyNCCL communicator	Gloo WORLD
tp_size>1 --disable-pynccl	NCCL	torch.distributed NCCL	另建 Gloo group

单卡仍会初始化 Gloo 进程组用于 _sync_get_memory 之类的通信，但不会创建 PyNCCL communicator ；只有默认多卡路径才会创建 PyNCCL communicator 。

# distributed/impl.py（节选）
class DistributedCommunicator:
    plugins: List[DistributedImpl] = [TorchDistributedImpl()]

    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
        return self.plugins[-1].all_reduce(x)

    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
        return self.plugins[-1].all_gather(x)

enable_pynccl_distributed会在 plugins 末尾追加 PyNCCLDistributedImpl，后续 all_reduce/all_gather 走最后一个 plugin。模型层只依赖统一的 DistributedCommunicator 接口，不关心底层是 PyNCCL 还是 torch.distributed。

---

3.3 权重、信息加载

并行层在构造时已经通过 get_tp_info 决定“本 rank 应装多大的参数”，models/weight.py 负责“从 checkpoint 取哪一片”。每个 rank 独立执行 load_weight，流式打开 safetensors，边读边按切分原则 先 shard、再按需 merge：

# models/weight.py（节选）
_SPLIT_DIM_0 = [".q_proj", ".k_proj", ".v_proj", ".gate_proj", ".up_proj"]
_SPLIT_DIM_1 = [".o_proj", ".down_proj"]

_MERGE_GROUPS = {...}
_SLOT_NAMES = {...}

# 按 TP 规则取出本 rank 对应的 shard（切分）
def _shard_tensor(key, value, r, n, num_kv_heads):
    if any(key.count(sub) for sub in _SPLIT_DIM_0):
        is_kv_proj = any(key.count(sub) for sub in (".k_proj", ".v_proj"))
        if is_kv_proj and num_kv_heads is not None and num_kv_heads < n:
            # GQA：按 head 切分，而非简单均分
            ...
            return value[head_idx * head_dim : (head_idx + 1) * head_dim].clone()
        return value.chunk(n, dim=0)[r].clone()
    elif any(key.count(sub) for sub in _SPLIT_DIM_1):
        return value.chunk(n, dim=1)[r].clone()
    elif key.count("lm_head") or key.count("embed_tokens"):
        # vocab 维按 rank 切分（div_ceil）
        ...
    else:
        # RMSNorm 等完整保留
        return value

# 按需 merge/stack（流式加载）
def load_weight(...):
    tp_info = get_tp_info()
    merge_buf: Dict[str, Dict[str, torch.Tensor]] = {}
    expert_buf: Dict[str, Dict[int, torch.Tensor]] = {}

    for file in tqdm(...):
        with safetensors.safe_open(...) as f:
            for name in f.keys():
                # 跳过多模态权重
                ...
                raw = f.get_tensor(name)
                name = name.removeprefix("language_model.")
                # 先按本 rank 切分
                tensor = _shard_tensor(...)
                del raw

                # 按需 merge（q/k/v → qkv_proj，gate/up → gate_up_proj）
                if (info := _get_merge_info(name)) is None:
                    out = (name, tensor)
                else:
                    # 按 merge 规则收集各 slot；未齐则 continue，齐了再拼接起来
                    ...
                    out = (merged_key, torch.cat(parts, dim=0))

                # MoE，把各 expert 同名权重（已 shard + merge）按 index stack 成专家维
                if config.is_moe and (expert_info := _get_expert_stack_info(out[0])) is not None:
                    # 收集各 expert；未齐则 continue，齐了再 stack 后 yield
                    # （此处是打包，不再做切分）
                    ...
                    yield packed_key, torch.stack(experts, dim=0)
                else:
                    # Dense 或已 merge 完成的普通权重，直接 yield
                    yield out[0], out[1]
    # 确认 checkpoint 是否完整加载
    assert not merge_buf, ...
    assert not expert_buf, ...

顺序是先 shard、再 merge：本 rank 已切好的 q/k/v 拼成 qkv_proj，gate/up 拼成 gate_up_proj。若先全局拼接再均分，会破坏本地 Attention / SwiGLU 的配对。MoE 模型还需要把各 expert 已处理的同名权重按 index “打包” 成单一 tensor 以及其他权重一起 yield；而 Dense 模型直接 yield 无需”打包“过程。

注：这种流式加载的目标是把峰值 CPU 内存控制在一个完整 tensor + 小 merge buffer 中，减少加载过程中的内存尖峰。

---

词表层在运行时的对应实现：

# layers/embedding.py（节选）
class VocabParallelEmbedding(BaseOP):
    def __init__(self, num_embeddings, embedding_dim):
        ...

    @nvtx_annotate("Embedding")
    def forward(self, x):
        y = indexing(weights=self.weight, indices=x,
                     vocab_range=self.vocab_range if self.tp_size > 1 else None)
        return self._comm.all_reduce(y) if self.tp_size > 1 else y


class ParallelLMHead(VocabParallelEmbedding):
    def load_state_dict(...):
        if not self.tied_embedding:
            return super().load_state_dict(...)
        # tie 时 pop 掉 lm_head.weight / bias（权重共享）
        ...

    def state_dict(...):
        ...

    @nvtx_annotate("LMHead")
    def forward(self, x):
        # Prefill阶段只取每个序列最后一个 token
        if batch.is_prefill:          
            x = x[batch.attn_metadata.get_last_indices(bs)].contiguous()
        module = self.tied_embedding or self
        logits = F.linear(x, module.weight, self.bias)
        if self.tp_size == 1:
            return logits
        output = self._comm.all_gather(logits)
        ...  # reshape 并截断到真实 num_embeddings
        return output[:, :self.num_embeddings]

• VocabParallelEmbedding：本地权重用 vocab_range 索引；tp_size > 1 时对输出做 all_reduce。
• ParallelLMHead：本地算 logits 后，tp_size > 1 时 all_gather 拼回完整词表 logits 并截断。

加载完成后，Engine 通过 self.model.load_state_dict(...) 把已 shard/merge 的 tensor 填入 meta 上建好的结构；然后还会完成 KV cache、page table、Attention/MoE backend、Sampler 以及 CUDA Graph 的初始化。

全部初始化结束后，系统才真正进入可服务状态，开始接收请求。

---

3.4 请求接收与广播

这个过程中只有 rank0 会和 Tokenizer / Detokenizer 交互，其他 rank 的请求消息全部来自 rank0 的广播，从而保证各 rank 看到相同的请求集合与顺序；单卡运行同理，只是没有跨 rank 同步：

# scheduler/io.py（节选）
def _recv_msg_single_rank(self, blocking: bool = False):
    pending_msgs: List[BaseBackendMsg] = []
    # 单卡：直接从 Tokenizer 拉取
    ...

# rank0：从 Tokenizer 收消息，用 Gloo 广播条数，用 ZMQ Pub 把内容发给其他 rank
def _recv_msg_multi_rank0(self, blocking: bool = False):
    pending_msgs: List[BaseBackendMsg] = []
    if blocking:
        # 收一条 raw，ZMQ Pub 发给其他 rank，再本地 decode
        ...

    pending_raw_msgs: List[bytes] = []
    while not self._recv_from_tokenizer.empty():
        pending_raw_msgs.append(self._recv_from_tokenizer.get_raw())

    # 仅用 Gloo 广播消息条数
    src_tensor = torch.tensor(len(pending_raw_msgs))
    self.tp_cpu_group.broadcast(src_tensor, root=0).wait()

    # 逐条 ZMQ Pub 发出，并在本地 decode
    for raw in pending_raw_msgs:
        self._send_into_ranks.put_raw(raw)
        pending_msgs.append(self._recv_from_tokenizer.decode(raw))
    return pending_msgs

def _recv_msg_multi_rank1(self, blocking: bool = False):
    pending_msgs: List[BaseBackendMsg] = []
    if blocking:
        self.run_when_idle()
        pending_msgs.append(self._recv_from_rank0.get())

    # 其他 rank 接收条数，再从 ZMQ Sub 取相同消息
    dst_tensor = torch.tensor(-1)
    self.tp_cpu_group.broadcast(dst_tensor, root=0).wait()
    for _ in range(int(dst_tensor.item())):
        pending_msgs.append(self._recv_from_rank0.get())
    return pending_msgs

def _reply_tokenizer_rank0(self, reply: List[DetokenizeMsg]):
    # 只有 rank0 回传 Detokenizer；其他 rank 的 _reply_tokenizer_rank1 为空操作
    ...

相同请求集合 + 相同请求顺序 + 相同 token/长度状态才构成可协同计算的 Batch，这是后续 RadixAttention 前缀缓存正确协同计算的前提。各 rank 根据相同消息独立更新本地 Prefill/Decode/Cache 状态，并执行相同的调度逻辑，运行时并没有每轮广播完整的 Batch 对象。

从启动接收请求到产出结果的完整路径，已在第 2 章一个请求的旅程中已经提到过了，此处就不再展开了。

4 对照真实 SGLang

前面的流程说明了 mini-SGLang 怎样启动多个 Scheduler、协调同一批请求，再通过权重分片与集合通信完成一次前向。真实 SGLang 的模块划分更细，但这些部分仍然可以大致对应：

功能	mini-SGLang（python/minisgl/）	SGLang（python/sglang/srt/）
多进程启动	server/launch.py	entrypoints/engine.py
调度与模型执行	scheduler/scheduler.py、engine/engine.py	managers/scheduler.py、managers/tp_worker.py、model_executor/model_runner
请求接收与协调	scheduler/io.py、scheduler/scheduler.py	managers/scheduler.py
通信组与张量通信	distributed/impl.py、kernel/pynccl.py	distributed/parallel_state.py
Column/Row/QKV Parallel	layers/linear.py、models/weight.py	layers/linear.py
词表切分与加载	layers/embedding.py、models/weight.py	layers/vocab_parallel_embedding.py

以上表格只列举部分模块仅供参考。SGLang 文件更多、配置与通信路径更复杂涉及 DP/PP/EP 等，但这条主线与 mini-SGLang 相通。可以先用 mini-SGLang 理解主线，再映射到 SGLang 的对应模块。

5 总结与测试题

5.1 课程总结

本章沿着 mini-SGLang 的源码，梳理了从多进程启动、按 TP rank 与切分原则加载权重，到接收并协调处理同一批请求的流程。多 rank 协作过程中的通信工具：Gloo同步信息条数、ZMQ广播消息内容、PyNCCL/NCCL前向中的张量集合通信共同保证了各 rank 状态一致以及最后计算结果的正确。

---

5.2 测试题

1. 每个 rank 都会启动一个 Scheduler，作用是什么？

提示：只考虑多卡运行的情况。

2. 权重是如何加载的？各 rank 又如何协调到同一批请求？

提示：分别考虑单卡与多卡。

3. 为什么采用流式加载权重，merge/stack 后再 yield 又是为了什么？

提示：关注加载过程中的峰值 CPU 内存，以及运行时层（如 fused qkv、MoE）所需的权重占用。

参考资料

• mini-SGLang
• SGLang
• 第 2 章一个请求的旅程